license: apache-2.0 language:

• en pipeline_tag: sentence-similarity inference: false

Monarch Mixer-BERT模型

这是一个80M参数的M2-BERT预训练模型检查点，序列长度达32768，并针对长上下文检索任务进行了微调。

查阅论文《Monarch Mixer: 一种基于GEMM的简易次二次方架构》以及我们的博客文章，了解我们如何训练这个长序列模型。

该模型由Jon Saad-Falcon、Dan Fu和Simran Arora训练完成。

访问我们的GitHub仓库获取下载和微调指南！

使用方法

可通过Hugging Face的AutoModel加载此模型：

from transformers import AutoModelForSequenceClassification
model = AutoModelForSequenceClassification.from_pretrained(
  "togethercomputer/m2-bert-80M-32k-retrieval",
  trust_remote_code=True
)

加载时会显示关于FlashFFTConv未使用参数的大段警告信息。如需加载带FlashFFTConv的模型，请参考我们的GitHub仓库。

该模型生成用于检索的嵌入向量，维度为768：

from transformers import AutoTokenizer, AutoModelForSequenceClassification

max_seq_length = 32768
testing_string = "每天早晨，我都会泡杯咖啡开始新的一天。"
model = AutoModelForSequenceClassification.from_pretrained(
  "togethercomputer/m2-bert-80M-32k-retrieval",
  trust_remote_code=True
)

tokenizer = AutoTokenizer.from_pretrained(
  "bert-base-uncased",
  model_max_length=max_seq_length
)
input_ids = tokenizer(
  [testing_string],
  return_tensors="pt",
  padding="max_length",
  return_token_type_ids=False,
  truncation=True,
  max_length=max_seq_length
)

outputs = model(**input_ids)
embeddings = outputs['sentence_embedding']

也可通过Together API获取嵌入向量（API密钥请见此处）：

import os
import requests

def generate_together_embeddings(text: str, model_api_string: str, api_key: str):
    url = "https://api.together.xyz/api/v1/embeddings"
    headers = {
        "accept": "application/json",
        "content-type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }
    session = requests.Session()
    response = session.post(
        url,
        headers=headers,
        json={
            "input": text,
            "model": model_api_string
        }
    )
    if response.status_code != 200:
        raise ValueError(f"请求失败，状态码{response.status_code}：{response.text}")
    return response.json()['data'][0]['embedding']

print(generate_together_embeddings(
  '你好世界',
  'togethercomputer/m2-bert-80M-32k-retrieval',
  os.environ['TOGETHER_API_KEY'])[:10]
)

致谢

Alycia Lee为AutoModel支持提供了帮助。

引用

若使用本模型或认为我们的工作有价值，请引用如下文献：

@inproceedings{fu2023monarch,
  title={Monarch Mixer: 一种基于GEMM的简易次二次方架构},
  author={傅丹毅 and 阿里拉·西蒙 and 格罗根·杰西卡 and 约翰逊·艾西斯 and 埃尤博格鲁·萨布里 and 托马斯·阿明·W and 斯佩克特·本杰明 and 波利·迈克尔 and 鲁德拉·阿特里 and 雷·克里斯托弗},
  booktitle={神经信息处理系统进展},
  year={2023}
}