PyTorch Geometric LLM 模块完全指南从知识图谱索引到图检索增强生成【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric导读torch_geometric.llm是 PyTorch Geometric 中面向「图神经网络 × 大语言模型」融合场景的专用子模块本文以其官方 API 文档 docs/source/modules/llm.rst 为骨架结合仓库源码、示例与测试系统讲解该模块的顶层基础设施大图索引与 RAG 查询、十余个开箱即用的 LLM 相关模型LLM 包装器、句子编码器、文本转知识图谱、LLM 评估器等以及 RAG 后端工具。读完本文你将掌握如何用LargeGraphIndexer把海量三元组压缩成唯一索引、如何用RAGQueryLoader搭建「查询 → 子图采样 → 特征装配」的检索流程以及如何跑通从原始文本到知识图谱再到 GNNLLM 问答的完整链路。模块总览llm 文档的 API 地图docs/source/modules/llm.rst是 Sphinx 自动生成文档的入口它通过autosummary与automodule指令将模块的公共 API 划分为三层模块顶层torch_geometric.llm直接导出的LargeGraphIndexer大图索引器与RAGQueryLoaderRAG 查询加载器见 torch_geometric/llm/init.py 中的classes列表Models 子模块torch_geometric.llm.models聚合了 10 个模型类Utils 子模块torch_geometric.llm.utils聚合了 3 个后端工具类。实际模块实现内容比文档骨架更丰富。__init__.py中的classes与models/__init__.py、utils/__init__.py共同勾勒出完整 API 清单层级公开类源码文件顶层LargeGraphIndexer、RAGQueryLoadertorch_geometric/llm/large_graph_indexer.py、torch_geometric/llm/rag_loader.pyModelsSentenceTransformer、VisionTransformer、LLM、LLMJudge、TXT2KG、GRetriever、MoleculeGPT、GLEM、ProteinMPNN、GITMoltorch_geometric/llm/models/UtilsKNNRAGFeatureStore、NeighborSamplingRAGGraphStore、DocumentRetrievertorch_geometric/llm/utils/模块定位清晰Models 解决「用什么模型编码文本/生成回答」Utils 解决「如何在远程后端上做检索与子图采样」顶层类则负责「索引与查询编排」。整套设计服务于把知识图谱当作 RAG 上下文喂给 LLM 的流水线。顶层基础设施LargeGraphIndexer与RAGQueryLoader用三元组构建唯一节点/边索引LargeGraphIndexer的核心思想是当一个数据集由多个子图组成而这些子图都属于一个更大的图时将子图中的节点与边按全局唯一 ID 汇总到一张索引表从而节省存储资源源码 docstring 原话collate the values into a large graph store to save resources。其构造参数与数据约定如下nodes: Iterable[str]图中全部节点 ID必须唯一否则抛出AttributeErroredges: KnowledgeGraphLike形如[(cats, eat, dogs)]的三元组序列(头实体, 关系, 尾实体)同样要求唯一node_attr/edge_attr可选的节点/边属性字典其中节点属性必须包含键pid节点 ID边属性必须包含e_pid、h、r、t、edge_idx五个键源码中NODE_KEYS与EDGE_KEYS常量定义了这些保留键。常用的构造入口是类方法from_triplets它从三元组流自动去重出节点集与边集还支持pre_transform回调对每个三元组做预处理from torch_geometric.llm import LargeGraphIndexer triplets [(cats, eat, dogs), (dogs, chase, cats)] indexer LargeGraphIndexer.from_triplets(triplets)collate类方法则可以把多个LargeGraphIndexer合并成一个全局唯一索引内部先各自转回三元组再统一重建。特征映射稀疏 ID 与稠密向量的桥梁索引器真正的价值在于特征管理。add_node_feature/add_edge_feature可以为每个唯一节点/边挂接新特征如文本编码后的向量当新特征与pid/e_pid不是一一对应、而是从某个已有特征映射而来时会用MappedFeature数据类包装并记录到_mapped_node_features/_mapped_edge_features集合中。get_node_features/get_edge_features支持按指定 PID 子集批量取特征且对 Tensor 类型特征走索引切片、对普通序列走迭代器兼顾效率与内存。持久化方面save(path)将节点/边索引pickle、映射标记pickle与全部属性torch.save按属性名分文件存于node_attr/、edge_attr/子目录落盘from_disk(path)反向还原加载时通过torch_geometric.io.fs.torch_load兼容多后端文件系统。to_data则把指定特征组装成标准Data对象含x、edge_index、edge_attr、node_id、edge_id。批量取特征get_features_for_triplets_groups模块级函数get_features_for_triplets_groups是索引器的高阶用法给定索引器与一组三元组分片如一个问答数据集为每个分片取出对应的节点特征、边特征与边索引批量构造成Data对象。其批处理策略值得注意用max_batch_size默认 250与num_workers默认取 CPU 核数动态计算理想批大小min(max_batch_size, max(1, len(triplet_groups) // num_workers))注释中明确了权衡批太小浪费算力每次取特征有固定开销批太大浪费内存需在内存中暂存全部嵌入通过ThreadPool并行取特征再按批切分拼接最终按Data对象逐个产出。RAG 查询协议与RAGQueryLoaderRAGQueryLoader面向远程后端上的 RAG 查询。它的输入是(feature_store, graph_store)二元组两者分别遵循 torch_geometric/llm/rag_loader.py 中定义的协议RAGFeatureStore协议retrieve_seed_nodes(query)查询与全部节点比对返回最接近的种子节点索引、retrieve_seed_edges(query)、load_subgraph(sample)把采样结果与特征装配成Data/HeteroData并提供config属性读写RAGGraphStore协议sample_subgraph(seed_nodes, seed_edges)基于种子做子图采样、register_feature_store(feature_store)异构图采样需要特征存储信息。RAGQueryLoader.query(query)的完整调用链rag_loader.py 第 133 行起为若配置了vector_retriever先用向量检索器取回相关文档retrieved_docs若augment_queryTrue把检索文档拼进查询query [query] retrieved_docs调用feature_store.retrieve_seed_nodes(query)得到种子节点调用graph_store.sample_subgraph(seed_nodes)采样子图调用feature_store.load_subgraph(sample...)装配特征若配置了subgraph_filter本地回调对结果做后处理若有向量检索器把data.text_context retrieved_docs附在结果上返回。config属性会在设置时通过_propagate_config同步下发给 feature store 与 graph store实现配置统一传播。Models 子模块开箱即用的 LLM 生态LLMHuggingFace 因果语言模型包装器LLMtorch_geometric/llm/models/llm.py是模块的核心其构造参数与默认值如下参数默认值说明model_name必填HuggingFace 模型名num_paramsNone模型参数量十亿用于自动估算 GPU 需求不指定时通过huggingface_hub.get_safetensors_metadata读取n_gpusNone手动指定 GPU 数量覆盖自动机制dtypetorch.bfloat16模型数据类型sys_promptNone系统提示词其 GPU 分配逻辑是一个基于启发式的自动决策required_memory 96.0 * num_params / 8.0源码注释说明约 8B 参数可放入 96GB GPU随后get_llm_kwargs依次探测各 GPU 空闲显存并累加直到满足需求若总显存不足则回退到纯 CPU 运行并发出警告建议用户显式传n_gpus修正。这是为规避 GPU OOM 而做的保守设计。初始化时还处理了 tokenizer 的兼容问题若模型带chat_template但无bos_token会用 dummy 对话通过apply_chat_template反推 BOSpad_token_id缺省置 0padding_side缺省为left。forward(question, answer, context, embedding)支持两种提示范式chat_template 范式推荐将sys_prompt与context - question组织成 system/user 消息经apply_chat_template(..., add_generation_promptTrue, enable_thinkingTrue)渲染再把 BOS 与可选的 RAG 嵌入拼接到inputs_embeds之前Llama 2 风格回退当模型无 chat_template 或无 sys_prompt 时走_get_embeds_old使用s[INST]、[/INST]、[/s]常量并warnings.warn提示改用带 chat template 的模型。回答标签经_label_input_ids截断到MAX_NEW_TOKENS128训练时用IGNORE_INDEX -100屏蔽非标签部分左填充保证 batch 对齐。inference则走generate并batch_decode输出文本。正因如此该包装器既能用于监督微调返回outputs.loss也能用于生成式推理。SentenceTransformer文本嵌入与多种池化策略SentenceTransformertorch_geometric/llm/models/sentence_transformer.py将任意 HuggingFace 编码器包装成句子嵌入器pooling_strategy支持四种PoolingStrategy枚举mean默认掩码加权平均、last最后一个 token兼容左填充的解码器、cls取首 token、last_hidden_state。输出统一做 L2 归一化F.normalize(emb, p2, dim1)。实现细节上max_seq_length会从config.max_position_embeddings与一次探针 token 化结果中取较小者避免某些模型配置缺失encode支持batch_size分块、output_device输出到指定设备遇到 GPU OOM 时自动回退 CPU 再切回原设备。TXT2KG文本到知识图谱TXT2KGtorch_geometric/llm/models/txt2kg.py把纯文本转化为(entity, relation, entity)三元组列表是「garbage in garbage out」链条的第一环。其核心参数参数默认值说明NVIDIA_NIM_MODELnvidia/llama-3.1-nemotron-70b-instructNIM 端点模型名NVIDIA_API_KEYNIM API 密钥为空且非本地模式时断言报错ENDPOINT_URLhttps://integrate.api.nvidia.com/v1自托管模型时的端点地址local_LMFalse使用本地 HuggingFace 模型开发调试用默认加载VAGOsolutions/SauerkrautLM-v2-14b-DPO并通过LLM.inference调用chunk_size512文本切块的最大字符数处理管线值得展开_chunk_text按句号/感叹号/问号等句末标点切块避免截断单词若块尾是字母则回退到最后一个空格云端模式下按进程数_get_num_procs取 CPU 亲和核数的一半把块等分用spawn上下文启动多进程starmap并行调用 NIMOpenAI 兼容接口temperature0, top_p1, max_tokens1024流式返回对known_reasoners如llama-3.1-nemotron-ultra-253b-v1、kimi-k2-instruct、gpt-oss等额外注入 detailed thinking on 系统消息调用失败按可重试性分类HTTP 429 或 ≥500 视为可重试指数退避 抖动重试单 worker 最多 200 次403 会提示检查 API 密钥整个多进程任务最多重试MAX_OUTER_RETRIES5次_parse_n_check_triples用纯 Python 规则清洗 LLM 输出容忍换行或括号两种格式、剥离引号、小写化、过滤空实体保证结果可被下游索引器直接消费_merge_triples_deterministically用str.casefold()做 Unicode 安全排序保证多进程结果可复现。add_doc_2_KG(txt, QA_pairNone)以 QA 对或自增文档 ID为键把三元组存入relevant_triples字典save_kg(path)用torch.save持久化。LLMJudgeLLM 评估器LLMJudgetorch_geometric/llm/models/llm_judge.py用 NIM 给(question, model_pred, correct_answer)三元组打分是评估 RAG/GNN 问答质量的工具。它内置两套评分提示词SYSTEM_PROMPT_1与SYSTEM_PROMPT_2均源自 NVIDIA Marlin Accuracy 方案要求模型只输出 4完全一致/ 2部分一致/ 0不一致。_process_score出于对 LLM 不守指令的容忍按 4/3/2/1/0 依次在响应中查找数字并归一化到 0–1score对两套提示各做多次重试分别最多 200 次与 20 次最后_average_scores取平均任一评分为 NaN 时退化为取最大值。领域专用模型与其余模型models/__init__.py还导出了六个面向特定领域的模型均有对应测试与示例GRetriever图检索增强的问答模型配套 examples/llm/g_retriever.py 与 test/llm/models/test_g_retriever.pyMoleculeGPT分子文本/图生成模型test/llm/models/test_molecule_gpt.pyGLEM图语言模型examples/llm/glem.pyProteinMPNN蛋白质序列逆折叠模型examples/llm/protein_mpnn.pyGITMol分子图-文本跨模态模型examples/llm/git_mol.pyVisionTransformer视觉 Transformer 编码器。Utils 子模块RAG 后端与向量检索工具DocumentRetriever轻量向量文档检索DocumentRetrievertorch_geometric/llm/utils/vectorrag.py实现VectorRetriever协议构造时传入raw_docs、可选的预嵌入embedded_docs、k_for_docs默认 2与编码模型若未提供嵌入则用编码器默认场景下是SentenceTransformer现算并把verbose参数从 kwargs 中弹出以免查询时刷屏。query对字符串查询先用torch.no_grad()编码再通过batch_knntorch_geometric/llm/utils/backend_utils.py 提供的批量最近邻返回 Top-k 原始文档。RAG 存储后端KNNRAGFeatureStore与NeighborSamplingRAGGraphStore这两个类分别实现RAGFeatureStore/RAGGraphStore协议构成「KNN 检索种子 邻居采样子图」的远程后端KNNRAGFeatureStore将查询编码后与全图节点特征做 KNN产出种子节点NeighborSamplingRAGGraphStore以种子节点为起点做邻居采样返回SamplerOutput供load_subgraph装配。实现细节见 torch_geometric/llm/utils/feature_store.py 与 torch_geometric/llm/utils/graph_store.py对应测试在 test/llm/utils/test_rag_feature_store.py 与 test/llm/utils/test_rag_graph_store.py。backend_utils图构建与 PCST 剪枝torch_geometric/llm/utils/backend_utils.py 提供流水线级的辅助函数preprocess_triplet三元组全部小写规范化create_graph_from_triples把三元组流变成图数据retrieval_via_pcst基于优先权 Steiner 树PCST依赖pcst_fast做图级检索剪枝——用余弦相似度给节点/边计算 prizetopk/topk_e控制保留数、cost_e控制边成本再运行pcst_fast选出与查询最相关的连通子图返回子图及文本化结果。对于缺特征的空图会直接跳过 PCSTcreate_remote_backend_from_graph_data把本地图数据转成远程后端LocalFeatureStoreLocalGraphStore。端到端实战跑通「文本 → KG → GNNLLM」问答链路仓库在 examples/llm/txt2kg_rag.py 提供了完整的参考实现把上述所有组件串成一条流水线。其架构阶段如下文本建图用TXT2KG把语料分块--kg_chunk_size默认 512 字符抽成三元组索引与远端化LargeGraphIndexer汇总三元组 →create_graph_from_triples建图 →create_remote_backend_from_graph_data生成远程后端查询RAGQueryLoader组合KNNRAGFeatureStoreNeighborSamplingRAGGraphStore可叠加DocumentRetriever做文档增强augment_query控制是否拼入检索文档GNN 编码子图特征送入GAT或SGFormer--gnn_model二选一得到图编码生成LLM以图嵌入为前缀embedding参数生成回答评估LLMJudge按 0–1 打分。关键命令行参数默认值均取自脚本顶部的常量定义参数默认值说明--gnn_modelGAT可选GAT/SGFormer--NV_NIM_MODELnvidia/llama-3.1-nemotron-ultra-253b-v1建图与评估用的 NIM 模型--NV_NIM_KEY必填NVIDIA API 密钥--llm_generator_namemeta-llama/Meta-Llama-3.1-8B-Instruct生成回答的 LLM--llm_generator_modefullfrozen/lora/full三种微调策略--gnn_hidden_channels1024GNN 隐层维度--num_gnn_layers4GNN 层数--lr/--epochs/--batch_size1e-5/2/1训练超参--wandb关闭可选 Weights Biases 日志该脚本还演示了LLM的frozen/lora/full三种生成器微调模式以及hf_hub_download下载预训练权重的用法。测试方面test/llm/test_large_graph_indexer.py、test/llm/test_rag_loader.py、test/llm/utils/test_rag_backend_utils.py 等覆盖了索引器往返、RAG 查询协议与后端工具的正确性可作为自定义实现的验证参照。适用前提与使用建议显存与模型规模LLM的 GPU 自动分配是启发式决策源码会在大显存不足时回退 CPU 并告警生产环境建议显式指定n_gpusNIM 依赖TXT2KG/LLMJudge云端模式依赖openai包与 NVIDIA NIM 端点temperature0、流式调用本地调试可用local_LMTrue但需能推理 14B 级模型源码注释明确提示小模型效果不可用HeteroGraph 兼容性源码多处标注 TODO: Make compatible with Heterographs目前 RAG 后端协议主要面向同构图依赖与安装pcst_fastPCST 剪枝、huggingface_hub、transformers均为可选依赖按需安装。总而言之torch_geometric.llm以「图数据为索引、图神经网络为编码器、LLM 为生成器」为设计主线提供了从数据建图、向量检索、子图采样到生成评估的一整套可组合构件是探索图检索增强生成Graph RAG路线的高质量工程化起点。【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考