Installation

conda env

• conda create --name mamba-aac python=3.12

install torch==2.7.0 first

• pip install torch==2.7.0

We recommend users to put mamba and casual-conv1d outside of the repository directory (to prevent nested .git detection in IDE)

• mamba

  • we use transformers's mamba implementation but it requires mamba package for faster inference
  • git clone https://github.com/state-spaces/mamba.git
  • git checkout tags/v2.2.4
  • pip install . --no-build-isolation
  • set assert False in modeling_mamba2.py in transformer package to disable calling reference (slow) function torch_forward()
  • It might take a long time depending on CPU speed.

• causal-conv1d

  • git clone https://github.com/Dao-AILab/causal-conv1d.git
  • git checkout tags/v1.5.0
  • pip install . --no-build-isolation

Install mamba_aac

• pip install -e .
• pip install timm==1.0.15 --no-deps

Training

• You need to modify few lines in modeling_mamba2.py huggingface's transformers library (4.54.1). read mamba_aac/models/custom_mamba.py and apply patches by yourself.
• Due to licensing restrictions, please download the datasets online and build your own JSON files based on your local directory path.

Inference

• train.py includes an inference routine. pass --generate and run the script with a single gpu.

Acknowledgement

Our implementation builds on the following works, and we thank the authors for their contributions.

• https://github.com/YuanGongND/ltu
• https://github.com/state-spaces/mamba
• https://github.com/Dao-AILab/causal-conv1d
• https://github.com/WenjunHuang94/ML-Mamba
• https://github.com/TRI-ML/prismatic-vlms
• https://github.com/huggingface/transformers
• https://github.com/cwx-worst-one/EAT