如何将Hugging Face推理端点生成的Embedding接入LlamaIndex节点？

问题 #

已通过Hugging Face推理端点实现Embedding生成，需将其接入LlamaIndex节点，复刻MetadataReplacementDemo的功能（原示例使用本地HuggingFaceEmbedding），用于生产环境。现有自定义Embedding生成代码，需替换官方示例中的本地嵌入模型配置。

实现方案 #

1. 适配LlamaIndex的EmbedModel接口 #

LlamaIndex的嵌入模型需实现embed_documents（批量文本嵌入）和embed_query（单条查询文本嵌入）两个核心方法，补充必要属性以兼容框架逻辑。修改自定义HuggingFaceEmbedder类：

endpoint = "你的Hugging Face推理端点URL"
api_key = "你的Hugging Face API密钥"

import requests

class HuggingFaceEmbedder:
    def __init__(self, endpoint, api_key):
        self.endpoint = endpoint
        self.api_key = api_key
        self.headers = {'Authorization': f'Bearer {self.api_key}'}
        # 补充模型名称属性（兼容LlamaIndex部分依赖逻辑）
        self.model_name = "sentence-transformers/all-mpnet-base-v2"

    def embed_documents(self, texts):
        if not isinstance(texts, list):
            raise ValueError("Input 'texts' must be a list of strings.")
        
        input_data = {"inputs": texts}
        res = requests.post(self.endpoint, headers=self.headers, json=input_data)
        
        if res.status_code == 200:
            return res.json()['embeddings']
        else:
            raise Exception(f"请求失败，状态码：{res.status_code}，响应：{res.text}")

    def embed_query(self, text):
        # 复用批量处理逻辑处理单条查询
        return self.embed_documents([text])[0]

2. 替换ServiceContext中的嵌入模型 #

将自定义HuggingFaceEmbedder实例传入ServiceContext，替代原有的本地HuggingFaceEmbedding：

from llama_index import ServiceContext, set_global_service_context
from llama_index.llms import OpenAI
from llama_index.node_parser import SentenceWindowNodeParser

# 初始化句子窗口节点解析器
node_parser = SentenceWindowNodeParser.from_defaults(
    window_size=3,
    window_metadata_key="window",
    original_text_metadata_key="original_text",
)

# 初始化LLM和自定义嵌入模型
llm = OpenAI(model="gpt-3.5-turbo", temperature=0.1)
embed_model = HuggingFaceEmbedder(endpoint, api_key)

# 创建全局服务上下文
ctx = ServiceContext.from_defaults(
    llm=llm,
    embed_model=embed_model,
    node_parser=node_parser,
)
set_global_service_context(ctx)

3. 后续流程与原示例保持一致 #

加载文档、生成节点、构建索引的逻辑完全复用原示例：

from llama_index import SimpleDirectoryReader, VectorStoreIndex

# 加载目标文档
documents = SimpleDirectoryReader(
    input_files=["./IPCC_AR6_WGII_Chapter03.pdf"]
).load_data()

# 生成带上下文窗口的节点
nodes = node_parser.get_nodes_from_documents(documents)

# 构建向量索引
sentence_index = VectorStoreIndex(nodes, service_context=ctx)

关键注意事项 #

• 维度一致性：确保Hugging Face推理端点的模型与原示例使用的sentence-transformers/all-mpnet-base-v2输出维度一致，否则会出现向量索引维度不匹配问题。
• 生产环境优化：建议给请求添加超时、重试机制，避免网络波动导致失败：

  from requests.adapters import HTTPAdapter
  from urllib3.util.retry import Retry
  
  class HuggingFaceEmbedder:
      def __init__(self, endpoint, api_key):
          self.endpoint = endpoint
          self.api_key = api_key
          self.headers = {'Authorization': f'Bearer {self.api_key}'}
          self.model_name = "sentence-transformers/all-mpnet-base-v2"
          # 初始化带重试的请求会话
          self.session = requests.Session()
          retry = Retry(total=3, backoff_factor=1, status_forcelist=[429, 500, 502, 503, 504])
          adapter = HTTPAdapter(max_retries=retry)
          self.session.mount("https://", adapter)
  
      def embed_documents(self, texts):
          if not isinstance(texts, list):
              raise ValueError("Input 'texts' must be a list of strings.")
          
          input_data = {"inputs": texts}
          try:
              res = self.session.post(self.endpoint, headers=self.headers, json=input_data, timeout=30)
              res.raise_for_status()
              return res.json()['embeddings']
          except requests.exceptions.RequestException as e:
              raise Exception(f"请求失败：{str(e)}")
  

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