库名称: transformers
许可证: mit
语言:

• 英语
  任务标签: 句子相似度
  标签:
• 文本嵌入
• 嵌入向量
• 信息检索
• beir
• 文本分类
• 语言模型
• 文本聚类
• 文本语义相似度
• 文本评估
• 文本重排序
• 特征提取
• 句子相似度
• 自然问题
• ms_marco
• 发烧
• hotpot_qa
• mteb

LLM2Vec: 大语言模型是潜在的强大文本编码器

LLM2Vec 是一种将仅解码器的大语言模型转换为文本编码器的简单方法。它包含三个简单步骤：1) 启用双向注意力，2) 掩码下一词预测，3) 无监督对比学习。该模型可以进一步微调以实现最先进的性能。

• 代码库: https://github.com/McGill-NLP/llm2vec
• 论文: https://arxiv.org/abs/2404.05961

安装

pip install llm2vec  

使用方式

from llm2vec import LLM2Vec  

import torch  
from transformers import AutoTokenizer, AutoModel, AutoConfig  
from peft import PeftModel  

# 加载基础 Mistral 模型，以及启用仅解码器 LLM 双向连接的自定义代码  
tokenizer = AutoTokenizer.from_pretrained(  
    "McGill-NLP/LLM2Vec-Meta-Llama-31-8B-Instruct-mntp"  
)  
config = AutoConfig.from_pretrained(  
    "McGill-NLP/LLM2Vec-Meta-Llama-31-8B-Instruct-mntp", trust_remote_code=True  
)  
model = AutoModel.from_pretrained(  
    "McGill-NLP/LLM2Vec-Meta-Llama-31-8B-Instruct-mntp",  
    trust_remote_code=True,  
    config=config,  
    torch_dtype=torch.bfloat16,  
    device_map="cuda" if torch.cuda.is_available() else "cpu",  
)  

# 加载 MNTP（掩码下一词预测）模型  
model = PeftModel.from_pretrained(  
    model,  
    "McGill-NLP/LLM2Vec-Meta-Llama-31-8B-Instruct-mntp",  
)  

# 用于编码和池化操作的封装器  
l2v = LLM2Vec(model, tokenizer, pooling_mode="mean", max_length=512)  

# 使用指令编码查询  
instruction = (  
    "给定一个网页搜索查询，检索回答该查询的相关段落："  
)  
queries = [  
    [instruction, "女性每天应摄入多少蛋白质"],  
    [instruction, "峰会定义"],  
]  
q_reps = l2v.encode(queries)  

# 编码文档。文档不需要指令  
documents = [  
    "作为一般指导原则，CDC 对 19 至 70 岁女性的蛋白质平均需求量为每天 46 克。但如本图表所示，如果您怀孕或正在为马拉松训练，则需要增加摄入量。查看下方图表了解每天应摄入的蛋白质量。",  
    "峰会定义（针对英语学习者）。1 山的最高点：山顶。2 最高级别。3 两国或多国政府领导人之间的会议或系列会议。",  
]  
d_reps = l2v.encode(documents)  

# 计算余弦相似度  
q_reps_norm = torch.nn.functional.normalize(q_reps, p=2, dim=1)  
d_reps_norm = torch.nn.functional.normalize(d_reps, p=2, dim=1)  
cos_sim = torch.mm(q_reps_norm, d_reps_norm.transpose(0, 1))  

print(cos_sim)  
"""  
tensor([[0.7724, 0.5563],  
        [0.4845, 0.5003]])  
"""  

问题

如果您对代码有任何疑问，请随时通过电子邮件联系 Parishad（parishad.behnamghader@mila.quebec）和 Vaibhav（vaibhav.adlakha@mila.quebec）。