南宁网站建设南宁,深圳做软件开发的公司,国外 网站设计,seo工作内容和薪资我们在前面的文章介绍了研究人员推出了一种挑战Transformer的新架构Mamba
他们的研究表明#xff0c;Mamba是一种状态空间模型(SSM)#xff0c;在不同的模式(如语言、音频和时间序列)中表现出卓越的性能。为了说明这一点#xff0c;研究人员使用Mamba-3B模型进行了语言建模…我们在前面的文章介绍了研究人员推出了一种挑战Transformer的新架构Mamba
他们的研究表明Mamba是一种状态空间模型(SSM)在不同的模式(如语言、音频和时间序列)中表现出卓越的性能。为了说明这一点研究人员使用Mamba-3B模型进行了语言建模实验。该模型超越了基于相同大小的Transformer的其他模型并且在预训练和下游评估期间它的表现与大小为其两倍的Transformer模型一样好。 Mamba的独特之处在于它的快速处理能力选择性SSM层以及受FlashAttention启发的硬件友好设计。这些特点使Mamba超越TransformerTransformer没有了传统的注意力和MLP块。
有很多人希望自己测试Mamba的效果所以本文整理了一个能够在Colab上完整运行Mamba代码代码中还使用了Mamba官方的3B模型来进行实际运行测试。
首先我们安装依赖这是官网介绍的 !pip install causal-conv1d1.0.0!pip install mamba-ssm1.0.1然后直接使用transformers库读取预训练的Mamba-3B import torchimport osfrom transformers import AutoTokenizerfrom mamba_ssm.models.mixer_seq_simple import MambaLMHeadModeltokenizer AutoTokenizer.from_pretrained(EleutherAI/gpt-neox-20b)model MambaLMHeadModel.from_pretrained(os.path.expanduser(state-spaces/mamba-2.8b), devicecuda, dtypetorch.bfloat16)可以看到3b的模型有11G 然后就是测试生成内容 tokens tokenizer(What is the meaning of life, return_tensorspt)input_ids tokens.input_ids.to(devicecuda)max_length input_ids.shape[1] 80fn lambda: model.generate(input_idsinput_ids, max_lengthmax_length, cgTrue,return_dict_in_generateTrue, output_scoresTrue,enable_timingFalse, temperature0.1, top_k10, top_p0.1,)out fn()print(tokenizer.decode(out[0][0]))这里还有一个chat的示例 import torchfrom transformers import AutoTokenizerfrom mamba_ssm.models.mixer_seq_simple import MambaLMHeadModeldevice cudatokenizer AutoTokenizer.from_pretrained(havenhq/mamba-chat)tokenizer.eos_token |endoftext|tokenizer.pad_token tokenizer.eos_tokentokenizer.chat_template AutoTokenizer.from_pretrained(HuggingFaceH4/zephyr-7b-beta).chat_templatemodel MambaLMHeadModel.from_pretrained(havenhq/mamba-chat, devicecuda, dtypetorch.float16)messages []user_message What is the date for announcementOn August 10 said that its arm JSW Neo Energy has agreed to buy a portfolio of 1753 mega watt renewable energy generation capacity from Mytrah Energy India Pvt Ltd for Rs 10,530 crore.messages.append(dict(roleuser,contentuser_message))input_ids tokenizer.apply_chat_template(messages, return_tensorspt, add_generation_promptTrue).to(cuda)out model.generate(input_idsinput_ids, max_length2000, temperature0.9, top_p0.7, eos_token_idtokenizer.eos_token_id)decoded tokenizer.batch_decode(out)messages.append(dict(roleassistant,contentdecoded[0].split(|assistant|\n)[-1]))print(Model:, decoded[0].split(|assistant|\n)[-1])这里我将所有代码整理成了Colab Notebook有兴趣的可以直接使用
https://avoid.overfit.cn/post/ed2d2cc2460d4e0683a270e2761e10ea
