如何在LangChain中用Llama-2隐藏状态构建嵌入模型并解决FAISS报错

我希望将Llama-2的hidden_states作为嵌入模型传入方法FAISS.from_documents(<filepath>, <embedding_model>)。目前已加载Llama-2模型，并能获取字符串的嵌入：

model_config = transformers.AutoConfig.from_pretrained(
    model_id,
    output_hidden_states=True,
    use_auth_token=auth_token,
)

# Load model directly
from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")

# Input data to test the code
input_text = "Hello World!"

encoded_input = tokenizer(input_text, return_tensors='pt')
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf",
                                            trust_remote_code=True,
                                            config=model_config,
                                            quantization_config=bnb_config,
                                            device_map='auto',
                                            use_auth_token=auth_token
                                            )

outputs = model(**encoded_input)
hidden_states = outputs.hidden_states

print(len(hidden_states))  # 33 for Llama-2: 1 (embeddings) + 32 (layers)
print(hidden_states[0].shape)  # Shape of the embeddings
print(hidden_states[2])

打印输出：

33
torch.Size([1, 4, 4096])
tensor([[[ 0.0373, -0.5762, -0.0180,  ...,  0.0962, -0.1099,  0.3767],
         [ 0.0676,  0.0400, -0.0033,  ...,  0.0655,  0.0278, -0.0079],
         [-0.0160,  0.0157,  0.0478,  ..., -0.0224, -0.0341,  0.0093],
         [ 0.0229, -0.0104,  0.0217,  ..., -0.0080, -0.0012, -0.0342]]],
       dtype=torch.float16, grad_fn=<ToCopyBackward0>)

尝试用Llama-2构建文档嵌入时：

from langchain.vectorstores import FAISS

# <clean> is the file-path
FAISS.from_documents(clean, model)

出现报错：

AttributeError: 'LlamaForCausalLM' object has no attribute 'embed_documents'

如何解决该问题，以及如何使用Llama-2的隐藏状态实现嵌入？

LangChain的FAISS.from_documents要求传入的嵌入模型必须实现embed_documents和embed_query方法，而LlamaForCausalLM是因果语言模型，本身没有这些方法，需自定义符合LangChain要求的嵌入类，基于Llama-2的hidden_states生成嵌入。

步骤1：自定义嵌入类 #

继承LangChain的Embeddings基类，实现embed_documents和embed_query方法，利用Llama-2的hidden_states生成文本表示，可选择最后几层隐藏状态的均值，或最后一个有效token的隐藏状态作为嵌入：

from langchain.embeddings.base import Embeddings
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

class LlamaHiddenStateEmbeddings(Embeddings):
    def __init__(self, model, tokenizer, hidden_layer_idx=-1, use_last_token=True):
        self.model = model
        self.tokenizer = tokenizer
        self.hidden_layer_idx = hidden_layer_idx  # 选择使用的隐藏层，-1代表最后一层
        self.use_last_token = use_last_token  # 是否用最后一个有效token的隐藏状态
        
        # 为Llama添加pad token（若无则用eos_token替代）
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token

    def _get_single_embedding(self, text):
        # 处理单条文本生成嵌入
        encoded_input = self.tokenizer(
            text,
            return_tensors='pt',
            padding=True,
            truncation=True,
            max_length=512
        ).to(self.model.device)
        
        # 禁用梯度计算提升速度
        with torch.no_grad():
            outputs = self.model(**encoded_input)
            target_hidden = outputs.hidden_states[self.hidden_layer_idx]
        
        if self.use_last_token:
            # 取最后一个有效token的隐藏状态（排除pad）
            attention_mask = encoded_input['attention_mask']
            last_token_pos = attention_mask.sum(dim=1) - 1
            embedding = target_hidden[torch.arange(target_hidden.shape[0]), last_token_pos, :]
        else:
            # 对所有有效token的隐藏状态取均值
            attention_mask = encoded_input['attention_mask'].unsqueeze(-1).expand(target_hidden.size())
            masked_hidden = target_hidden * attention_mask
            embedding = torch.sum(masked_hidden, dim=1) / torch.clamp(attention_mask.sum(dim=1), min=1e-9)
        
        return embedding.cpu().numpy()[0]

    def embed_documents(self, texts):
        # 批量处理文档生成嵌入
        return [self._get_single_embedding(text) for text in texts]

    def embed_query(self, text):
        # 处理查询文本生成嵌入
        return self._get_single_embedding(text)

步骤2：初始化嵌入模型并构建FAISS向量库 #

用已加载的Llama-2模型和tokenizer初始化自定义嵌入类，再传入FAISS.from_documents：

# 初始化自定义嵌入模型，可调整hidden_layer_idx和use_last_token参数
llama_embeddings = LlamaHiddenStateEmbeddings(model, tokenizer, hidden_layer_idx=-1, use_last_token=True)

# 构建FAISS向量库
db = FAISS.from_documents(clean, llama_embeddings)

关键细节 #

• 隐藏层选择：可根据任务效果调整hidden_layer_idx，比如用倒数第二层（-2）或多个层的均值，不同层的语义侧重不同。
• 嵌入方式选择：
  • 最后一个有效token的嵌入：适合短句，保留结尾语义重点。
  • token均值嵌入：适合长文本，捕获整体语义。
• 性能优化：若处理大量文档，可修改embed_documents实现批量输入，减少模型调用次数，提升处理效率。

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