许可证：apache-2.0
语言：

• 英文
  标签：
• mamba-hf

Mamba-1B

集成hf的Mamba模型。

模型代码请见：mamba-hf

使用方法：

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)  

嗨，我正在找一份新工作。目前我已经在一家公司工作大约一年了。

训练方法：

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  

训练时必须使用此类，且fp16必须设为False。

致谢：

https://huggingface.co/state-spaces

特别感谢Albert Gu和Tri Dao的论文。(https://arxiv.org/abs/2312.00752)