本文将简单介绍如何使用 AutoSchemaKG 提升 RAG(Retrieval Augmented Generation,检索增强生成) 的问答效果。实现的主要思路是:
- 解析知识图谱文件:用 NetworkX 解析 GraphML 知识图谱文件;
- 嵌入:用 FAISS (Facebook AI Similarity Search) 矢量化节点(实体、事件、概念),并将边的 “源节点 + 关系 + 目标节点” 拼接为字符串之后再做嵌入;
- 语义检索:
- 检索 “源节点 + 关系 + 目标节点” 的矢量库,找出最相似的文档,再用 大语言模型(LLM)对结果进一步筛选,筛选出节点;
- 检索相似的 passage(文本段落);
- 将上述包含分数的结果合并,将分数作为初始权重,通过 NetworkX 的 Pagerank 方法在 知识图谱 中通过节点间的关联关系传播权重,权重更新后,选择权重最高的 passage(文本段落) 为结果。
PageRank 核心原理(知识图谱场景):
- 给节点赋予初始权重(与查询相关的节点权重高)
- 每个节点将自身权重按 “出边数量” 平均分配给相邻节点(例如,节点 A 有 2 条出边,就将自身权重的 1/2 分别传给两个邻居)
- 迭代传播,直到权重分布稳定(误差小于 tol),最终每个节点的 pr[node] 即为 “综合相关度分数”。
通过权重传播,与查询节点间接关联的节点(例如 “苹果手机”→“苹果公司”→“芯片供应商”)也能获得较高分数,挖掘出 “字面不直接匹配但语义相关” 的内容。
传统的 RAG(Retrieval Augmented Generation,检索增强生成) 一般是直接查询返回相似的 passage(文本段落),没有上述通过知识图谱调整权重的过程。
代码实现
下面的代码使用 AutoSchemaKG 实现了 语义检索:
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com' # 设置 HuggingFace 镜像源,加快下载
os.environ["TRANSFORMERS_PROGRESS_BAR"] = "1" # 启用下载进度条
from sentence_transformers import SentenceTransformer
from atlas_rag.vectorstore.embedding_model import NvEmbed, SentenceEmbedding
# 使用中文通用语义嵌入模型 BAAI/bge-m3
encoder_model_name = "BAAI/bge-m3"
sentence_model = SentenceTransformer(
encoder_model_name,
trust_remote_code=True,
cache_folder="./model", # 模型会下载到这个目录,第一次执行会下载模型,比较慢。
model_kwargs={'device_map': "auto"}
)
sentence_encoder = SentenceEmbedding(sentence_model)
from openai import OpenAI
from atlas_rag.llm_generator import LLMGenerator
from configparser import ConfigParser
# Load OpenRouter API key from config file
config = ConfigParser()
config.read('config.ini')
reader_model_name = "qwen3"
client = OpenAI(
base_url="http://127.0.0.1:11434/v1",
api_key='ollama',
)
llm_generator = LLMGenerator(client=client, model_name=reader_model_name)
from atlas_rag.vectorstore import create_embeddings_and_index
keyword = 'law'
current_file_path = os.path.abspath(__file__)
current_dir = os.path.dirname(current_file_path)
working_directory = os.path.join(current_dir, f'target-big/{keyword}')
data = create_embeddings_and_index(
sentence_encoder=sentence_encoder,
model_name = encoder_model_name,
working_directory=working_directory,
keyword=keyword,
include_concept=True,
include_events=True,
normalize_embeddings= True,
text_batch_size=64,
node_and_edge_batch_size=64,
)
from atlas_rag.retriever import HippoRAG2Retriever
from atlas_rag import setup_logger
hipporag2_retriever = HippoRAG2Retriever(
llm_generator=llm_generator,
sentence_encoder=sentence_encoder,
data = data,
)
上述代码首先加载 GraphML 生成了索引以及相关文件,如果 GraphML 比较大,生成的过程会比较漫长,所以生成后会自动存储,避免未来启动时再重复生成这些文件。
我使用了部署在本地的 ollama 中的 qwen3 8b,由于 OpanAI 几乎成了所有大模型的规范,所以如果想使用线上的其它大模型,只需要修改一下 base_url 和 api_key 就能直接使用了。
测试效果
使用下面的代码提问:
# perform retrieval
content, sorted_context_ids = hipporag2_retriever.retrieve("黑龙江对出省印制的出版物需要经过哪个部门批准?", topN=3)
print(f"Retrieved content: {content}")
# start benchmarking
sorted_context = "\n".join(content)
result = llm_generator.generate_with_context("黑龙江对出省印制的出版物需要经过哪个部门批准?", sorted_context, max_new_tokens=2048, temperature=0.5)
print(result)
Retrieved content: ['第十九条 出版单位实行编辑责任制度。编审人员应当恪尽职守,保证出版物的质量。 第二十条 任何单位和个人不得以任何名义直接或间接购买或租借许可证、书号、刊号和版号,并参与出版、印刷或复制、发行等活动。 任何单位和个人不得伪造、假冒出版单位或报纸、期刊名称出版出版物。 第二十一条 出版物刊登的广告内容,应当真实、健康、科学、准确,符合广告法律、法规的规定。出版单位不得刊登非法、虚假广告,不得以新闻报道形式刊登广告,不得擅自更改广告审批机关审核的广告内容。 广告专版、专刊、专页、专栏应当有明显的广告标识。 第二十
二条 出版单位采集、编辑、发表新闻或组织、编辑、印发稿件,不得向报道对象和供稿者索取和收受审稿、编辑、印发等费用。 ...。']
hought: 根据第二十八条,省外出版物在本省印制需经省出版行政部门核发准印证。问题中“出省印制”指黑龙江出版物到省外印制,需经黑龙江省出版行政部门批准,再由受委托地省级部门核发准印
证。因此,黑龙江对出省印制的出版物需经省出版行政部门批准。
Answer: 省出版行政部门
总结
能让 AutoSchemaKG 在中文环境下跑起来并不是特别容易,实现思路也比单纯靠矢量数据库进行语义检索要复杂很多。
不过从理论上来说:这种 通过知识图谱提升RAG能力 的办法应该是可行的。
🪐感谢您观看,祝好运🪐