如何从Hugging Face的meta-llama/Llama-2-13b-chat-hf获取句嵌入？

Llama-2-13b-chat-hf是为对话生成优化的模型，并非专门的句嵌入模型，直接用AutoModel提取默认token的输出会不符合预期。以下是几种可行的解决方法：

1. 调整嵌入提取逻辑，适配模型特性 #

不要直接依赖<s>或</s>这类特殊token的输出，改用最后一层隐藏状态的均值（过滤掉padding token）来生成句嵌入，同时注意输入格式要符合Llama-2 Chat的要求：

from transformers import AutoTokenizer, AutoModel
import torch

model_name = "meta-llama/Llama-2-13b-chat-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name, torch_dtype=torch.bfloat16)

def get_sentence_embedding(sentence):
    # 构造适配语义搜索的输入格式
    input_text = f"<s>[INST] 请表示这句话的语义用于搜索：{sentence} [/INST]"
    inputs = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=512)
    
    with torch.no_grad():
        outputs = model(**inputs)
    
    # 计算有效token的隐藏状态均值
    mask = inputs["attention_mask"].unsqueeze(-1).expand(outputs.last_hidden_state.size())
    sum_embeddings = torch.sum(outputs.last_hidden_state * mask, dim=1)
    avg_embeddings = sum_embeddings / torch.clamp(mask.sum(dim=1), min=1e-9)
    # L2归一化提升语义匹配效果
    avg_embeddings = torch.nn.functional.normalize(avg_embeddings, p=2, dim=1)
    
    return avg_embeddings.cpu().numpy()[0]

# 测试示例
sample_sentence = "Llama-2模型的句嵌入提取方法"
embedding = get_sentence_embedding(sample_sentence)
print(f"嵌入维度：{embedding.shape}")  # 输出 (4096,)

2. 参考SGPT的非对称语义搜索思路 #

SGPT通过给查询和文档添加不同前缀引导模型生成适配搜索的嵌入，你可以直接套用这个逻辑：

• 查询文本前缀："Query: " + 你的句子
• 待匹配文本前缀："Passage: " + 你的句子

修改上述代码的input_text部分即可实现，这种方式能让模型针对语义搜索任务生成更精准的嵌入。

3. 使用专门微调的句嵌入版本 #

如果不想自己调整逻辑，可以直接使用基于Llama-2-13b-chat微调的句嵌入模型，这类模型已针对语义匹配任务优化，能直接输出可用的句嵌入。

内容的提问来源于stack exchange，提问作者Mukesh Reddy