LlamaIndex 在 Intel Gaudi 上的嵌入集成：GaudiEmbedding 安装、参数解析与 Graph RAG 实战-编程实验室

LlamaIndex 在 Intel Gaudi 上的嵌入集成：GaudiEmbedding 安装、参数解析与 Graph RAG 实战

【免费下载链接】llama_indexLlamaIndex is the document processing platform for AI项目地址: https://gitcode.com/GitHub_Trending/ll/llama_index

本篇技术指南聚焦 LlamaIndex 官方集成包llama-index-embeddings-gaudi，完整讲解如何在 Intel Gaudi（Habana）加速器上运行 HuggingFace 嵌入模型：从安装依赖、初始化GaudiEmbedding、理解其全部构造参数与底层 tokenize 覆盖逻辑，到以thenlper/gte-large为例完成文本/查询嵌入，并进一步跑通基于 Neo4j 知识图谱与本地 Gaudi LLM 的 Graph RAG 完整链路。读完本文，你将掌握在 Gaudi 硬件上构建 LlamaIndex 嵌入层与知识图谱检索应用的直接可运行方案。

一、集成包概览：什么是 llama-index-embeddings-gaudi

llama-index-embeddings-gaudi是 LlamaIndex 官方为 Intel Gaudi 加速器（HPU）提供的嵌入模型接入层，其核心类是llama_index.embeddings.gaudi.GaudiEmbedding（源码位于 llama_index/embeddings/gaudi/base.py）。它通过包装 HuggingFace 生态的SentenceTransformer模型，让原本面向 CPU/GPU 的文本嵌入流程可以在 Gaudi HPU 上获得加速执行。

从包配置 pyproject.toml 可以确认该包的版本与依赖约束：

当前版本：0.4.0；
Python 要求：>=3.10,<4.0；
运行时依赖：optimum[habana]>=1.21.2与llama-index-core>=0.13.0,<0.15。

其中optimum[habana]是 Intel 官方在 HuggingFace Optimum 框架上的 Habana 适配分支，负责把 Transformer/SentenceTransformer 模型调度到 HPU 上执行，是运行本集成的硬件驱动基础。从源码结构看，GaudiEmbedding直接继承llama_index.core.base.embeddings.base.BaseEmbedding，因此它与 LlamaIndex 的索引构建、检索器、查询引擎天然兼容，可以无缝替换默认的OpenAIEmbedding等实现。

二、环境准备与安装

2.1 硬件前提

本集成面向 Intel Gaudi 加速器（如 Gaudi 1/2 及后续型号）设计，示例脚本通过--device hpu参数显式指定推理设备。运行前需要确认宿主机已正确安装 Intel Gaudi 软件栈（含 HPU 驱动与torch的 Habana 版本）。

2.2 安装命令

在 Intel Gaudi 主机上按以下顺序安装（摘自 examples/README.md）：

pip install --upgrade-strategy eager optimum[habana] pip install llama-index-embeddings-gaudi

第一条命令安装带 Habana 适配的 Optimum 生态（使用--upgrade-strategy eager是为了让相关依赖统一升级到与 Habana 软件栈匹配的版本）；第二条命令安装本集成包本身。

2.3 安装验证

安装完成后，可通过导入与实例化快速验证（对应 examples/basic.py 的写法）：

from llama_index.embeddings.gaudi import GaudiEmbedding embed_model = GaudiEmbedding( embedding_input_size=-1, model_name="thenlper/gte-large", )

若没有抛出 ImportError 且模型加载完成，说明环境就绪。GaudiEmbedding由包入口 llama_index/embeddings/gaudi/init.py 导出，导入路径为llama_index.embeddings.gaudi。

三、基础用法：文本嵌入与查询嵌入

3.1 最小可运行示例

以下代码来自 examples/basic.py，演示了 GaudiEmbedding 的最基本调用方式：

from llama_index.embeddings.gaudi import GaudiEmbedding if __name__ == "__main__": embed_model = GaudiEmbedding( embedding_input_size=-1, model_name="thenlper/gte-large", ) # Basic embedding example embeddings = embed_model.get_text_embedding("It is raining cats and dogs here!") print(len(embeddings), embeddings[:10])

运行方式（需在 Gaudi 主机上，启用 HPU 惰性执行与分布式集合模式）：

PT_HPU_LAZY_ACC_PAR_MODE=1 PT_HPU_ENABLE_LAZY_COLLECTIVES=true python basic.py

PT_HPU_LAZY_ACC_PAR_MODE=1与PT_HPU_ENABLE_LAZY_COLLECTIVES=true是 Habana 运行时的环境开关，前者控制惰性累积的并行执行模式，后者启用惰性集合通信，是官方示例要求的启动前提。

3.2 面向 LlamaIndex 的完整 API

GaudiEmbedding作为BaseEmbedding的子类，向 LlamaIndex 上层暴露以下嵌入接口（见 base.py）：

方法	说明
`get_text_embedding(text)`	对单条文本生成向量，底层调用`_get_text_embedding`
`get_text_embeddings(texts)`	批量生成向量，底层调用`_get_text_embeddings`
`get_query_embedding(query)`	对查询语句生成向量，底层调用`_get_query_embedding`
`aget_query_embedding(query)`	查询嵌入的异步版本
`aget_text_embedding(text)`	文本嵌入的异步版本

所有方法最终汇聚到_embed()方法，它调用底层GaudiSentenceTransformer.encode()完成真正的向量化，并返回 Python 列表格式的浮点向量：

def _embed(self, sentences, prompt_name=None): return self._model.encode( sentences, batch_size=self.embed_batch_size, prompt_name=prompt_name, normalize_embeddings=self.normalize, ).tolist()

3.3 在索引/查询流程中接入

因为GaudiEmbedding实现了 LlamaIndex 标准嵌入接口，可以直接挂到全局Settings.embed_model，或在构建索引、查询引擎时作为embed_model参数传入（下文 Graph RAG 示例即采用传入参数的方式）。向量维度由所选模型决定，例如thenlper/gte-large输出 1024 维向量，具体可通过len(embeddings)验证。

四、GaudiEmbedding 构造参数与源码级解析

4.1 构造参数一览

根据 base.py 的__init__签名，GaudiEmbedding支持以下参数：

参数	类型	默认值	说明
`model_name`	`str`	`"thenlper/gte-large"`	HuggingFace 模型名或本地模型路径
`embedding_input_size`	`int`	`-1`	传给底层 transformer 的输入长度控制（见 4.2）
`max_length`	`Optional[int]`	`512`	输入最大长度（`gt=0`约束），与`embedding_input_size`配合决定 tokenize 的`max_length`
`normalize`	`bool`	`True`	是否对输出向量做 L2 归一化，直接透传给`encode(normalize_embeddings=...)`
`query_instruction`	`Optional[str]`	`None`	拼接在查询文本前的指令（当前实现中该字段保留但未激活 prompts 映射，见 4.3）
`text_instruction`	`Optional[str]`	`None`	拼接在文本前的指令（同上）
`tokenizer`	`Optional[Any]`	`None`	预留的自定义 tokenizer 参数
`embed_batch_size`	`int`	`DEFAULT_EMBED_BATCH_SIZE`	批量嵌入的批次大小，透传给`encode(batch_size=...)`
`callback_manager`	`Optional[CallbackManager]`	`None`	LlamaIndex 回调管理器，用于链路追踪
`**model_kwargs`	—	—	透传给`SentenceTransformer`的其余关键字（如`cache_folder`等）

需要注意：示例代码中的embedding_input_size并不在 pydantic 字段声明中，而是通过**model_kwargs透传给底层的GaudiSentenceTransformer。

4.2 核心实现：GaudiSentenceTransformer 的 tokenize 覆盖

GaudiSentenceTransformer是SentenceTransformer的子类（base.py），其关键改动是重写了tokenize方法，以适配 HPU 上静态 shape 的执行特性：

def tokenize(self, texts): return self._first_module().tokenizer( texts, max_length=self.max_seq_length if ( self.embedding_input_size == -1 or self.embedding_input_size > self.max_seq_length ) else self.embedding_input_size, padding="max_length", return_tensors="pt", truncation=True, )

该逻辑的含义：

当embedding_input_size == -1时，使用模型自身的max_seq_length作为max_length；
当embedding_input_size大于max_seq_length时，同样退回到max_seq_length（避免超出模型能力）；
其他情况下，使用用户指定的embedding_input_size作为max_length；
统一采用padding="max_length"（补齐到定长）与truncation=True（超长截断），从而保证 HPU 上每个 batch 的输入 shape 固定，减少重编译开销。

这解释了为何示例统一传入embedding_input_size=-1：让模型按自身最大序列长度处理输入，同时保持张量形状稳定。

4.3 指令（instruction）字段的当前实现状态

从源码看，query_instruction与text_instruction字段在构造时被记录（base.py），但原本用于把指令注入SentenceTransformerprompts 的代码块当前处于注释状态，且_embed调用时prompt_name=None。因此可以推断：在当前版本中，指令字段主要用于保持与 HuggingFace 集成包 API 的一致性，指令文本暂不会自动拼接到输入。若需为 BGE/Instructor 类模型附加检索指令，可参考同目录 utils.py 中预留的指令模板（DEFAULT_EMBED_INSTRUCTION、DEFAULT_QUERY_INSTRUCTION、DEFAULT_QUERY_BGE_INSTRUCTION_EN/ZH以及get_query_instruct_for_model_name/get_text_instruct_for_model_name辅助函数），自行在预处理阶段拼接。

4.4 默认模型与模型缓存

包内 utils.py 定义了默认模型常量为"BAAI/bge-small-en-v1.5"，而 base.py 中GaudiEmbedding的默认模型为"thenlper/gte-large"。模型加载时通过cache_folder=get_cache_dir()（即 LlamaIndex 的统一缓存目录）复用已下载的 HuggingFace 权重，避免重复下载。utils.py同时保留了 BGE 系列与 Instructor 系列模型的清单及对应指令生成逻辑，可作为选择其他模型的参考。

五、Graph RAG 实战：Gaudi 本地 LLM + Neo4j 知识图谱

官方示例 graphrag.py 展示了在 Intel Gaudi 上完全本地化运行 Graph RAG（图分析与检索增强生成）的完整方案：用GaudiEmbedding生成节点嵌入、用GaudiLLM做三元组抽取与回答生成、用 Neo4j 存储知识图谱，最后通过图查询完成检索问答。

5.1 架构与数据流

从 graphrag.py 的run_code()可以还原整体流程：

用WikipediaReader加载维基百科页面文本（示例为 "Guardians of the Galaxy Vol. 3"）；
初始化GaudiLLM（模型HuggingFaceH4/zephyr-7b-alpha，配置messages_to_prompt与query_wrapper_prompt）与GaudiEmbedding（模型thenlper/gte-large）；
通过Settings.llm = llm、Settings.chunk_size = 512配置全局默认值；
连接 Neo4j（Neo4jGraphStore），构建StorageContext；
用KnowledgeGraphIndex.from_documents(...)从文档抽取三元组并生成知识图谱索引（max_triplets_per_chunk=3，include_embeddings=True）；
用index.as_query_engine(...)（embedding_mode="hybrid"、similarity_top_k=5、response_mode="tree_summarize"）执行图检索问答；
将查询结果按 zephyr 对话模板包装后打印答案。

该示例同时依赖 LlamaIndex 的其他官方集成包：llama-index-llms-gaudi（Gaudi LLM）、llama-index-graph-stores-neo4j（Neo4j 图存储）、llama-index-readers-wikipedia（维基百科读取器），详见 examples/requirements.txt。

5.2 第一步：启动 Neo4j 数据库服务器

Graph RAG 示例需要 Neo4j（含 APOC 插件）作为图存储后端。官方给出的 Docker 启动命令如下：

docker run --restart always --publish=7474:7474 --publish=7687:7687 --env NEO4J_AUTH=neo4j/<neo4j-server-password> -v $PWD/data:/data -v $PWD/plugins:/plugins --name neo4j-apoc -e NEO4J_apoc_export_file_enabled=true -e NEO4J_apoc_import_file_enabled=true -e NEO4J_apoc_import_file_use__neo4j__config=true -e NEO4JLABS_PLUGINS=\[\"apoc\"\] -e NEO4J_dbms_security_procedures_unrestricted=apoc.\\\* neo4j:5.22.0

要点解读：

--publish=7474:7474暴露 Neo4j 浏览器界面（HTTP），--publish=7687:7687暴露 Bolt 协议端口供驱动连接；
NEO4J_AUTH=neo4j/<password>设置初始账号密码（需替换<neo4j-server-password>）；
-v $PWD/data:/data与-v $PWD/plugins:/plugins挂载数据与插件目录；
通过多个NEO4J_apoc_*环境变量启用 APOC 的导入导出能力，并安装apoc插件、放开过程执行限制；
镜像版本固定为neo4j:5.22.0。

5.3 第二步：安装附加依赖并设置环境变量

官方文档要求 Intel Gaudi 软件版本1.18.0 或更高，然后安装：

pip install llama-index-llms-huggingface pip install llama-index-llms-gaudi pip install requirements.txt

其中requirements.txt即本包的 examples/requirements.txt，包含llama-index-graph-stores-neo4j、llama-index-readers-wikipedia、wikipedia、InstructorEmbedding==1.0.1、python-dotenv等。

随后设置 Neo4j 连接环境变量（对应 graphrag.py 中的读取逻辑）：

export NEO4J_USERNAME=neo4j export NEO4J_PASSWORD=<neo4j-server-password> #default: neo4j export NEO4J_URL=neo4j://<neo4j-server-host-ip>:7687 export NEO4J_DATABASE=neo4j

四个变量分别对应Neo4jGraphStore的username、password、url（Bolt 地址）与database，缺一不可，否则图存储连接会失败。

5.4 第四步：运行 Graph RAG 示例

PT_HPU_LAZY_ACC_PAR_MODE=1 PT_HPU_ENABLE_LAZY_COLLECTIVES=true python graphrag.py

运行时会先构建知识图谱（LLM 抽取三元组 + 嵌入生成节点向量），随后对问题"List the cast of Guardians of the Galaxy Vol. 3"进行图检索（混合嵌入模式）与tree_summarize式回答合成。示例 graphrag.py 的setup_parser()还内置了大量 Gaudi 推理调优参数（如--bf16、--use_hpu_graphs、--max_new_tokens、--batch_size、--use_flash_attention、--bucket_size等），可通过命令行覆盖默认值，用于吞吐与延迟调优。

六、常见问题与注意事项

必须在 Gaudi 硬件上运行：optimum[habana]与 HPU 惰性执行环境变量是运行前提，普通 CPU/GPU 环境无法发挥该集成价值；
输入长度控制：需要控制嵌入输入长度时设置embedding_input_size，建议保持-1（使用模型默认max_seq_length）以简化 shape 管理；
向量归一化：normalize=True（默认）会对输出做 L2 归一化，适合余弦相似度检索场景；若底层向量库内部已归一化，可关闭以减少计算；
Graph RAG 的软件版本约束：官方明确 Intel Gaudi 软件版本需 ≥ 1.18.0，且 Neo4j 使用5.22.0镜像，升级前应核对版本兼容性；
指令字段暂不自动生效：当前版本的query_instruction/text_instruction仅作为字段保留，需要指令注入时请自行在输入文本前拼接（参考 utils.py 中的模板）。