最新消息（2023年5月）：请迁移至 e5-small-v2，该版本性能更优且使用方法保持不变。
通过弱监督对比预训练实现文本嵌入。 Liang Wang、Nan Yang、Xiaolong Huang、Binxing Jiao、Linjun Yang、Daxin Jiang、Rangan Majumder、Furu Wei，arXiv 2022
该模型包含12层，嵌入维度为384。
以下是对 MS-MARCO 段落排序数据集中的查询和段落进行编码的示例。
from openmind import AutoTokenizer, AutoModel, is_torch_npu_available
from openmind_hub import snapshot_download
import torch.nn.functional as F
from torch import Tensor
import argparse
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument(
"--model_name_or_path",
type=str,
help="Path to model",
default="zhouhui/e5-small",
)
args = parser.parse_args()
return args
def average_pool(last_hidden_states: Tensor,
attention_mask: Tensor) -> Tensor:
last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
def main():
args = parse_args()
model_path = args.model_name_or_path
if is_torch_npu_available():
device = "npu:0"
else:
device = "cpu"
# Each input text should start with "query: " or "passage: ".
# For tasks other than retrieval, you can simply use the "query: " prefix.
input_texts = ['query: how much protein should a female eat',
'query: summit define',
"passage: As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
"passage: Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments."]
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModel.from_pretrained(model_path)
# Tokenize the input texts
batch_dict = tokenizer(input_texts, max_length=512, padding=True, truncation=True, return_tensors='pt')
outputs = model(**batch_dict)
embeddings = average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
# normalize embeddings
embeddings = F.normalize(embeddings, p=2, dim=1)
scores = (embeddings[:2] @ embeddings[2:].T) * 100
print(scores.tolist())
if __name__ == "__main__":
main()请参考我们的论文：https://arxiv.org/pdf/2212.03533.pdf。
查看 unilm/e5 以复现 BEIR 和 MTEB 基准 上的评估结果。
以下是使用 sentence_transformers 的示例。
from sentence_transformers import SentenceTransformer
model = SentenceTransformer('intfloat/e5-small')
input_texts = [
'query: how much protein should a female eat',
'query: summit define',
"passage: As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
"passage: Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments."
]
embeddings = model.encode(input_texts, normalize_embeddings=True)pip install sentence_transformers~=2.2.2
贡献者：michaelfeil
1. 是否需要为输入文本添加前缀 "query: " 和 "passage: "？
是的，这是模型的训练方式，否则会导致性能下降。
以下是一些经验法则：
对于非对称任务（如开放域问答中的段落检索、特定主题信息检索），请分别使用 "query: " 和 "passage: "。
对于对称任务（如语义相似度、复述检索），使用 "query: " 前缀。
如果希望将嵌入用作特征（如线性探针分类、聚类），使用 "query: " 前缀。
2. 为什么我复现的结果与模型卡片中报告的结果略有不同？
transformers 和 pytorch 的不同版本可能会导致微小但非零的性能差异。
3. 为什么余弦相似度分数分布在 0.7 到 1.0 之间？
这是已知且预期的行为，因为我们在 InfoNCE 对比损失中使用了 0.01 的低温度参数。
对于文本检索或语义相似度等文本嵌入任务，重要的是分数的相对顺序而非绝对值，因此这不应构成问题。
如果您发现我们的论文或模型有帮助，请考虑按以下方式引用：
@article{wang2022text,
title={Text Embeddings by Weakly-Supervised Contrastive Pre-training},
author={Wang, Liang and Yang, Nan and Huang, Xiaolong and Jiao, Binxing and Yang, Linjun and Jiang, Daxin and Majumder, Rangan and Wei, Furu},
journal={arXiv preprint arXiv:2212.03533},
year={2022}
}此模型仅适用于英文文本。长文本将被截断为最多 512 个 token。