🚀 Mistral-SUPRA

Mistral-SUPRA 模型基于 Mistral-7B 初始化，经过进一步训练转化为线性 RNN。它是论文《Linearizing Large Language Models》的配套模型，能在推理时兼具 Transformer 和循环模型的功能。

🚀 快速开始

要使用 Mistral-SUPRA 模型，需先安装支持线性注意力的 OpenLM 分支：

pip install git+https://github.com/tri-ml/linear_open_lm.git

导入 OpenLM 类：

from open_lm.open_lm_hf import *

使用 AutoTokenizer 和 AutoModelForCausalLM 加载模型：

from open_lm.open_lm_hf import *
from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("tri-ml/mistral-supra")
model = AutoModelForCausalLM.from_pretrained("tri-ml/mistral-supra")

inputs = tokenizer(["Machine learning is"], return_tensors="pt")
gen_kwargs = {"max_new_tokens": 50, "top_p": 0.8, "temperature": 0.8, "do_sample": True, "repetition_penalty": 1.1}
output = model.generate(inputs['input_ids'], **gen_kwargs)
output = tokenizer.decode(output[0].tolist(), skip_special_tokens=True)
print(output)
# Machine learning is a branch of artificial intelligence (AI) that enables computers to learn from experience without being explicitly programmed. Machine learning is used in a wide range of applications, including spam filtering, image recognition, speech recognition, and computer-based medical diagnosis

模型支持并行和循环两种模式：

# 循环模式
output = model.to('cuda').generate(inputs['input_ids'].to('cuda'), use_cache=True, **gen_kwargs)

# 并行模式
output = model.to('cuda').generate(inputs['input_ids'].to('cuda'), use_cache=False, **gen_kwargs)

✨ 主要特性

• 基于 Mistral-7B 模型初始化，转化为线性模型，兼具 Transformer 和循环模型的功能。
• 可在推理时根据 use_cache 参数选择并行或循环模式。

📦 安装指南

要使用该模型，需先安装支持线性注意力的 OpenLM 分支：

pip install git+https://github.com/tri-ml/linear_open_lm.git

💻 使用示例

基础用法

from open_lm.open_lm_hf import *
from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("tri-ml/mistral-supra")
model = AutoModelForCausalLM.from_pretrained("tri-ml/mistral-supra")

inputs = tokenizer(["Machine learning is"], return_tensors="pt")
gen_kwargs = {"max_new_tokens": 50, "top_p": 0.8, "temperature": 0.8, "do_sample": True, "repetition_penalty": 1.1}
output = model.generate(inputs['input_ids'], **gen_kwargs)
output = tokenizer.decode(output[0].tolist(), skip_special_tokens=True)
print(output)
# Machine learning is a branch of artificial intelligence (AI) that enables computers to learn from experience without being explicitly programmed. Machine learning is used in a wide range of applications, including spam filtering, image recognition, speech recognition, and computer-based medical diagnosis

高级用法

# 循环模式
output = model.to('cuda').generate(inputs['input_ids'].to('cuda'), use_cache=True, **gen_kwargs)

# 并行模式
output = model.to('cuda').generate(inputs['input_ids'].to('cuda'), use_cache=False, **gen_kwargs)

📚 详细文档

模型详情

• 开发者：Toyota Research Institute
• 模型类型：这是一个自回归语言模型，基于 Mistral-7B 初始化，并根据 SUPRA 架构训练为线性模型。
• 数据集：基于 Mistral-7B 初始化，在 100B 个 RefinedWeb 标记上进行训练。
• 分词器：mistralai/Mistral-7B-v0.1
• 库：OpenLM（使用支持线性注意力的 分支）
• 许可证：该模型遵循 Apache License, Version 2.0 许可。

训练详情

• Mistral-SUPRA 在 128 个 H100 80GB GPU 上使用 AWS SageMaker 进行训练。
• 对 100B 个标记的训练在 1.5 天内完成。

性能评估

使用 Eleuther LM Eval Harness 仓库进行评估。

以下是 Mistral-SUPRA 与其他类似规模模型的性能对比：

🔧 技术细节

该模型基于 Mistral-7B 初始化，通过特定的训练过程将其转化为线性 RNN。其线性注意力代码可在 https://github.com/TRI-ML/linear_open_lm/ 找到。在推理时，可根据 use_cache 参数选择并行或循环模式，以满足不同的应用需求。

📄 许可证

该模型遵循 Apache License, Version 2.0 许可。

📚 引用说明

如果使用该模型，请引用论文《Linearizing Large Language Models》：

@article{Mercat2024Linearizing,
  title={Linearizing Large Language Models},
  author={Jean Mercat and Igor Vasiljevic and Sedrick Keh and Kushal Arora and Achal Dave and Adrien Gaidon and Thomas Kollar},
  year={2024},
  journal={arXiv preprint arXiv:2405.06640},
}

OpenLM 引用

@misc{open_lm,
  author = {Gururangan, Suchin and Wortsman, Mitchell and Gadre, Samir Yitzhak and Dave, Achal and Kilian, Maciej and Shi, Weijia and Mercat, Jean and Smyrnis, Georgios and Ilharco, Gabriel and Jordan, Matt and Heckel, Reinhard and Dimakis, Alex and Farhadi, Ali and Shankar, Vaishaal and Schmidt, Ludwig},
  title = {{open_lm}:  a minimal but performative language modeling (LM) repository},
  year = {2023},
  note = {GitHub repository},
  url = {https://github.com/mlfoundations/open_lm/}
}