Mamba 1B

🚀 Mamba-1B

Mamba-1B是集成了Hugging Face（hf_integration）的模型，可用於相關的自然語言處理任務，為用戶提供便捷的建模體驗。

🚀 快速開始

模型使用

以下是使用Mamba-1B模型進行文本生成的示例代碼：

from transformers import AutoModelForCausalLM , AutoTokenizer

model = AutoModelForCausalLM.from_pretrained('Q-bert/Mamba-1B', trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained('Q-bert/Mamba-1B')

text = "Hi"

input_ids = tokenizer.encode(text, return_tensors="pt")

output = model.generate(input_ids, max_length=20, num_beams=5, no_repeat_ngram_size=2)

generated_text = tokenizer.decode(output[0], skip_special_tokens=True)

print(generated_text)

運行上述代碼後，可能的輸出示例如下：

Hi, I'm looking for a new job. I've been working at a company for about a year now.

模型訓練

以下是使用自定義訓練類進行Mamba-1B模型訓練的示例代碼：

from transformers import Trainer ,TrainingArguments
import torch
import os


class MambaTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        input_ids = inputs.pop("input_ids")
        lm_logits = model(input_ids)[0]

        labels = input_ids.to(lm_logits.device)
        shift_logits = lm_logits[:, :-1, :].contiguous()
        labels = labels[:, 1:].contiguous()

        loss_fct = torch.nn.CrossEntropyLoss()
        lm_loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), labels.view(-1))

        return lm_loss

⚠️ 重要提示

訓練時必須使用上述自定義的MambaTrainer類，並且fp16必須設置為False。

📄 許可證

本項目採用Apache-2.0許可證。

🔗 參考資料

• 模型相關信息可參考：mamba-hf
• 原模型相關資源：https://huggingface.co/state-spaces
• 特別感謝Albert Gu和Tri Dao的文章：(https://arxiv.org/abs/2312.00752)