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

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

LLM2Vec 是一种将仅解码器架构的大语言模型（LLM）转换为文本编码器的简易方案。该方案包含三个简单步骤：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模型及支持双向连接的定制代码。MNTP LoRA权重已合并至基础模型  
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",  
)  
model = PeftModel.from_pretrained(  
    model,  
    "McGill-NLP/LLM2Vec-Meta-Llama-31-8B-Instruct-mntp",  
)  
model = model.merge_and_unload()  # CPU上可能需要数分钟  

# 加载无监督SimCSE模型（在MNTP模型基础上加载训练好的LoRA权重）  
model = PeftModel.from_pretrained(  
    model, "McGill-NLP/LLM2Vec-Meta-Llama-31-8B-Instruct-mntp-unsup-simcse"  
)  

# 创建编码与池化操作封装器  
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.6007, 0.3518],  
        [0.4131, 0.4855]])  
"""  

问题咨询

如有疑问请联系Parishad（parishad.behnamghader@mila.quebec）或Vaibhav（vaibhav.adlakha@mila.quebec）