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

• 英文
  任务标签: 文本相似度
  标签:
• 文本嵌入
• 嵌入向量
• 信息检索
• beir
• 文本分类
• 语言模型
• 文本聚类
• 文本语义相似度
• 文本评估
• 文本重排序
• 特征提取
• 句子相似度
• 自然问题
• 微软宏数据集
• 事实验证
• 热点问答
• 多任务评估基准

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模型及支持双向连接的定制代码  
tokenizer = AutoTokenizer.from_pretrained(  
    "McGill-NLP/LLM2Vec-Sheared-LLaMA-mntp"  
)  
config = AutoConfig.from_pretrained(  
    "McGill-NLP/LLM2Vec-Sheared-LLaMA-mntp", trust_remote_code=True  
)  
model = AutoModel.from_pretrained(  
    "McGill-NLP/LLM2Vec-Sheared-LLaMA-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-Sheared-LLaMA-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.8180, 0.5825],  
        [0.1069, 0.1931]])  
"""  

问题咨询

如有任何疑问，请联系Parishad（parishad.behnamghader@mila.quebec）和Vaibhav（vaibhav.adlakha@mila.quebec）。