知識圖譜檢索增強的GraphRAG(基于Neo4j代碼實現(xiàn)）
作者：hanscalZheng
為了存儲知識圖譜數(shù)據(jù)，首先需要搭建一個 Neo4j 實例。最簡單的方法是在 Neo4j Aura 上啟動一個免費實例，它提供了 Neo4j 數(shù)據(jù)庫的云版本。當然，也可以通過docker本地啟動一個，然后將圖譜數(shù)據(jù)導(dǎo)入到Neo4j 數(shù)據(jù)庫中。
圖檢索增強生成（Graph RAG）正逐漸流行起來，成為傳統(tǒng)向量搜索方法的有力補充。這種方法利用圖數(shù)據(jù)庫的結(jié)構(gòu)化特性，將數(shù)據(jù)以節(jié)點和關(guān)系的形式組織起來，從而增強檢索信息的深度和上下文關(guān)聯(lián)性。圖在表示和存儲多樣化且相互關(guān)聯(lián)的信息方面具有天然優(yōu)勢，能夠輕松捕捉不同數(shù)據(jù)類型間的復(fù)雜關(guān)系和屬性。而向量數(shù)據(jù)庫在處理這類結(jié)構(gòu)化信息時則顯得力不從心，它們更擅長通過高維向量處理非結(jié)構(gòu)化數(shù)據(jù)。在 RAG 應(yīng)用中，結(jié)合結(jié)構(gòu)化的圖數(shù)據(jù)和非結(jié)構(gòu)化的文本向量搜索，可以讓我們同時享受兩者的優(yōu)勢，這也是本文將要探討的內(nèi)容。
構(gòu)建知識圖譜通常是利用圖數(shù)據(jù)表示的強大功能中最困難的一步。它需要收集和整理數(shù)據(jù)，這需要對領(lǐng)域知識和圖建模有深刻的理解。為了簡化這一過程，可以參考已有的項目或者利用LLM來創(chuàng)建知識圖譜，進而可以把重點放在檢索召回，以及LLM的生成階段。下面來進行相關(guān)代碼的實踐。
1.知識圖譜構(gòu)建
為了存儲知識圖譜數(shù)據(jù)，首先需要搭建一個 Neo4j 實例。最簡單的方法是在 Neo4j Aura 上啟動一個免費實例，它提供了 Neo4j 數(shù)據(jù)庫的云版本。當然，也可以通過docker本地啟動一個，然后將圖譜數(shù)據(jù)導(dǎo)入到Neo4j 數(shù)據(jù)庫中。
步驟I：Neo4j環(huán)境搭建
下面是本地啟動docker的運行示例：
docker run -d \
--restart always \
--publish=7474:7474 --publish=7687:7687 \
--env NEO4J_AUTH=neo4j/000000 \
--volume=/yourdockerVolume/neo4j:/data \
neo4j:5.19.0步驟II：圖譜數(shù)據(jù)導(dǎo)入
演示中，我們可以使用伊麗莎白一世的維基百科頁面。利用 LangChain 加載器從維基百科獲取并分割文檔，后存入Neo4j數(shù)據(jù)庫。為了試驗中文上的效果，我們導(dǎo)入這個Github上的這個項目（QASystemOnMedicalKG）中的醫(yī)學(xué)知識圖譜，包含近35000個節(jié)點，30萬組三元組，大致得到如下結(jié)果：
圖片
或者利用LangChainLangChain 加載器從維基百科獲取并分割文檔，大致如下面步驟所示：
# 讀取維基百科文章
raw_documents = WikipediaLoader(query="Elizabeth I").load()
# 定義分塊策略
text_splitter = TokenTextSplitter(chunk_size=512, chunk_overlap=24)
documents = text_splitter.split_documents(raw_documents[:3])
llm=ChatOpenAI(temperature=0, model_name="gpt-4-0125-preview")
llm_transformer = LLMGraphTransformer(llm=llm)
# 提取圖數(shù)據(jù)
graph_documents = llm_transformer.convert_to_graph_documents(documents)
# 存儲到 neo4j
graph.add_graph_documents(
graph_documents,
baseEntityLabel=True,
include_source=True
)2.知識圖譜檢索
在對知識圖譜檢索之前，需要對實體和相關(guān)屬性進行向量嵌入并存儲到Neo4j數(shù)據(jù)庫中：
- 實體信息向量嵌入：將實體名稱和實體的描述信息拼接后，利用向量表征模型進行向量嵌入（如下述示例代碼中的add_embeddings方法所示）。
- 圖譜結(jié)構(gòu)化檢索：圖譜的結(jié)構(gòu)化檢索分為四個步驟：步驟一，從圖譜中檢索與查詢相關(guān)的實體；步驟二，從全局索引中檢索得到實體的標簽；步驟三，根據(jù)實體標簽在相應(yīng)的節(jié)點中查詢鄰居節(jié)點路徑；步驟四，對關(guān)系進行篩選，保持多樣性（整個檢索過程如下述示例代碼中的structured_retriever方法所示）。
class GraphRag(object):
def __init__(self):
"""Any embedding function implementing `langchain.embeddings.base.Embeddings` interface."""
self._database = 'neo4j'
self.label = 'Med'
self._driver = neo4j.GraphDatabase.driver(
uri=os.environ["NEO4J_URI"],
auth=(os.environ["NEO4J_USERNAME"],
os.environ["NEO4J_PASSWORD"]))
self.embeddings_zh = HuggingFaceEmbeddings(model_name=os.environ["EMBEDDING_MODEL"])
self.vectstore = Neo4jVector(embedding=self.embeddings_zh,
username=os.environ["NEO4J_USERNAME"],
password=os.environ["NEO4J_PASSWORD"],
url=os.environ["NEO4J_URI"],
node_label=self.label,
index_name="vector"
)
def query(self, query: str, params: dict = {}) -> List[Dict[str, Any]]:
"""Query Neo4j database."""
from neo4j.exceptions import CypherSyntaxError
with self._driver.session(database=self._database) as session:
try:
data = session.run(query, params)
return [r.data() for r in data]
except CypherSyntaxError as e:
raise ValueError(f"Generated Cypher Statement is not valid\n{e}")
def add_embeddings(self):
"""Add embeddings to Neo4j database."""
# 查詢圖中所有節(jié)點，并且根據(jù)節(jié)點的描述和名字生成embedding，添加到該節(jié)點上
query = """MATCH (n) WHERE not (n:{}) RETURN ID(n) AS id, labels(n) as labels, n""".format(self.label)
print("qurey node...")
data = self.query(query)
ids, texts, embeddings, metas = [], [], [], []
for row in tqdm(data,desc="parsing node"):
ids.append(row['id'])
text = row['n'].get('name','') + row['n'].get('desc','')
texts.append(text)
metas.append({"label": row['labels'], "context": text})
self.embeddings_zh.multi_process = False
print("node embeddings...")
embeddings = self.embeddings_zh.embed_documents(texts)
print("adding node embeddings...")
ids_ret = self.vectstore.add_embeddings(
ids=ids,
texts=texts,
embeddings=embeddings,
metadatas=metas
)
return ids_ret
# Fulltext index query
def structured_retriever(self, query, limit=3, simlarity=0.9) -> str:
"""
Collects the neighborhood of entities mentioned in the question
"""
# step1 從圖譜中檢索與查詢相關(guān)的實體。
docs_with_score = self.vectstore.similarity_search_with_score(query, k=topk)
entities = [item[0].page_content for item in data if item[1] > simlarity] # score
self.vectstore.query(
"CREATE FULLTEXT INDEX entity IF NOT EXISTS FOR (e:Med) ON EACH [e.context]")
result = ""
for entity in entities:
qry = entity
# step2 從全局索引中查出entity label，
query1 = f"""CALL db.index.fulltext.queryNodes('entity', '{qry}') YIELD node, score
return node.label as label,node.context as context, node.id as id, score LIMIT {limit}"""
data1 = self.vectstore.query(query1)
# step3 根據(jù)label在相應(yīng)的節(jié)點中查詢鄰居節(jié)點路徑
for item in data1:
node_type = item['label']
node_type = item['label'] if type(node_type) == str else node_type[0]
node_id = item['id']
query2 = f"""match (node:{node_type})-[r]-(neighbor) where ID(node) = {node_id} RETURN type(r) as rel, node.name+' - '+type(r)+' - '+neighbor.name as output limit 50"""
data2 = self.vectstore.query(query2)
# step4 為了保持多樣性，對關(guān)系進行篩選
rel_dict = defaultdict(list)
if len(data2) > 3*limit:
for item1 in data2:
rel_dict[item1['rel']].append(item1['output'])
if rel_dict:
rel_dict = {k:random.sample(v, 3) if len(v)>3 else v for k,v in rel_dict.items()}
result += "\n".join(['\n'.join(el) for el in rel_dict.values()]) +'\n'
else:
result += "\n".join([el['output'] for el in data2]) +'\n'
return result3.結(jié)合LLM生成
最后利用大語言模型(LLM)根據(jù)從知識圖譜中檢索出來的結(jié)構(gòu)化信息，生成最終的回復(fù)。下面的代碼中我們以通義千問開源的大語言模型為例：
步驟I：加載LLM模型
from langchain import HuggingFacePipeline
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
def custom_model(model_name, branch_name=None, cache_dir=None, temperature=0, top_p=1, max_new_tokens=512, stream=False):
tokenizer = AutoTokenizer.from_pretrained(model_name,
revision=branch_name,
cache_dir=cache_dir)
model = AutoModelForCausalLM.from_pretrained(model_name,
device_map='auto',
torch_dtype=torch.float16,
revision=branch_name,
cache_dir=cache_dir
)
pipe = pipeline("text-generation",
model = model,
tokenizer = tokenizer,
torch_dtype = torch.bfloat16,
device_map = 'auto',
max_new_tokens = max_new_tokens,
do_sample = True
)
llm = HuggingFacePipeline(pipeline = pipe,
model_kwargs = {"temperature":temperature, "top_p":top_p,
"tokenizer":tokenizer, "model":model})
return llm
tongyi_model = "Qwen1.5-7B-Chat"
llm_model = custom_model(model_name=tongyi_model)
tokenizer = llm_model.model_kwargs['tokenizer']
model = llm_model.model_kwargs['model']步驟II：輸入檢索數(shù)據(jù)生成回復(fù)
final_data = self.get_retrieval_data(query)
prompt = ("請結(jié)合以下信息，簡潔和專業(yè)的來回答用戶的問題，若信息與問題關(guān)聯(lián)緊密，請盡量參考已知信息。\n"
"已知相關(guān)信息:\n{context} 請回答以下問題:{question}".format(cnotallow=final_data, questinotallow=query))
messages = [
{"role": "system", "content": "你是**開發(fā)的智能助手。"},
{"role": "user", "content": prompt}
]
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
model_inputs = tokenizer([text], return_tensors="pt").to(self.device)
generated_ids = model.generate(model_inputs.input_ids,max_new_tokens=512)
generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)]
response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
print(response)4 結(jié)語
對一個查詢問題分別進行了測試, 與沒有RAG僅利用LLM生成回復(fù)的的情況進行對比，在有GraphRAG 的情況下，LLM模型回答的信息量更大、準確會更高。
責(zé)任編輯：武曉燕
來源：
AI大模型咨訊
