Mamba-Chat 🐍

Mamba-Chat 是第一个基于状态空间模型架构而非 Transformer 的聊天语言模型。

该模型基于 Albert Gu 和 Tri Dao 的研究成果《Mamba：具有选择性状态空间的线性时间序列建模》（论文）以及他们的模型实现。本仓库提供了基于 Huggingface Trainer 类的一些修改而来的模型训练/微调代码。

Mamba-Chat 基于 Mamba-2.8B，并在 HuggingFaceH4/ultrachat_200k 数据集的 16,000 个样本上进行了微调。想了解更多，你可以：

• 在 Huggingface 上查看模型 🤗
• 在 Haven 社区 Discord 上与我们交流 🧑‍🤝‍🧑
• 在 Google Colab 上与 Mamba-Chat 对话

运行 Mamba-Chat

我们提供了测试和微调我们模型的代码。以下是如何开始使用以及你可以用它做什么：

克隆仓库并安装依赖：

git clone https://github.com/havenhq/mamba-chat.git
cd mamba-chat
pip install -r requirements.txt

与 Mamba-Chat 对话（命令行聊天机器人）：

python chat.py

与 Mamba-Chat 对话（gradio 应用）：

pip install gradio==4.8.0
python app.py --share

在 Ultrachat 数据集的子集上微调 Mamba（基础模型）：

python train_mamba.py --model state-spaces/mamba-2.8b --tokenizer EleutherAI/gpt-neox-20b --learning_rate 5e-5 --batch_size 4 --data_path ./data/ultrachat_small.jsonl --num_epochs 3

如果你有 24GB 显存的显卡（3090、4090 等），可以使用以下设置：

python train_mamba.py --model state-spaces/mamba-2.8b --tokenizer EleutherAI/gpt-neox-20b --learning_rate 5e-5 --batch_size 1 --gradient_accumulation_steps 4 --optim paged_adamw_8bit --data_path ./data/ultrachat_small.jsonl --num_epochs 3

引用

bibtex
@misc{haven2023mambachat,
  title        = {Mamba-Chat},
  author       = {Justus Mattern and Konstantin Hohr},
  year         = {2023},
  howpublished = {GitHub},
  url          = {https://github.com/havenhq/mamba-chat}
}