利用Milvus向量數(shù)據(jù)庫,帶你實(shí)現(xiàn)GraphRAG 原創(chuàng)
教你如何結(jié)合Milvus向量數(shù)據(jù)庫實(shí)現(xiàn)GraphRAG技術(shù)。
GraphRAG技術(shù)借助知識圖譜,給RAG應(yīng)用注入了新的動(dòng)力,使其能夠在海量數(shù)據(jù)中精確檢索所需信息。本文將帶你了解GraphRAG的實(shí)現(xiàn)方法,包括如何創(chuàng)建索引以及如何利用Milvus向量數(shù)據(jù)庫進(jìn)行查詢,助你在信息檢索的道路上事半功倍。
1、先決條件
在運(yùn)行本文中的代碼之前,請確保已安裝以下依賴項(xiàng):
pip install --upgrade pymilvus
pip install git+https://github.com/zc277584121/graphrag.git
注意:通過一個(gè)分支倉庫來安裝GraphRAG,這是因?yàn)镸ilvus的存儲(chǔ)功能在本文編寫時(shí)還未被官方正式合并。
2、數(shù)據(jù)準(zhǔn)備
為了進(jìn)行GraphRAG索引,我們需要準(zhǔn)備一個(gè)小型文本文件。我們將從Gutenberg項(xiàng)目(https://www.gutenberg.org/)下載一個(gè)大約一千行的文本文件,這個(gè)文件包含了關(guān)于達(dá)芬奇的故事。
利用這個(gè)數(shù)據(jù)集,構(gòu)建一個(gè)涉及達(dá)芬奇所有關(guān)系的知識圖譜索引,并使用Milvus向量數(shù)據(jù)庫來檢索相關(guān)知識,以便回答相關(guān)問題。
以下是Python代碼,用于下載文本文件并進(jìn)行初步處理:
import nest_asyncio
nest_asyncio.apply()
import os
import urllib.request
index_root = os.path.join(os.getcwd(), 'graphrag_index')
os.makedirs(os.path.join(index_root, 'input'), exist_ok=True)
url = "https://www.gutenberg.org/cache/epub/7785/pg7785.txt"
file_path = os.path.join(index_root, 'input', 'davinci.txt')
urllib.request.urlretrieve(url, file_path)
with open(file_path, 'r+', encoding='utf-8') as file:
# 使用文本文件的前934行,因?yàn)楹竺娴男信c本例無關(guān)。
# 如果想節(jié)省API密鑰成本,可以截?cái)辔谋疚募詼p小大小。
lines = file.readlines()
file.seek(0)
file.writelines(lines[:934]) # 如果想節(jié)省API密鑰成本,可以減少這個(gè)數(shù)字。
file.truncate()
3、初始化工作空間
現(xiàn)在,使用GraphRAG對文本文件進(jìn)行索引。首先運(yùn)行??graphrag.index --init?
?命令初始化工作空間。
python -m graphrag.index --init --root ./graphrag_index
4、配置環(huán)境變量文件
在索引的根目錄下,能找到一個(gè)名為??.env?
?的文件。要啟用這個(gè)文件,請將你的OpenAI API密鑰添加進(jìn)去。
注意事項(xiàng):
- 本例將使用OpenAI模型作為一部分,請準(zhǔn)備好你的API密鑰。
- GraphRAG索引的成本相對較高,因?yàn)樗枰肔LM處理整個(gè)文本語料庫。運(yùn)行這個(gè)演示可能會(huì)花費(fèi)一些資金。為了節(jié)省成本,你可以考慮將文本文件縮減尺寸。
5、執(zhí)行索引流程
運(yùn)行索引需要一些時(shí)間,請耐心等待。執(zhí)行完畢后,你可以在??./graphrag_index/output/<timestamp>/?
?路徑下找到一個(gè)新創(chuàng)建的文件夾,里面包含了多個(gè)parquet格式的文件。
執(zhí)行以下命令開始索引過程:
python -m graphrag.index --root ./graphrag_index
6、使用Milvus向量數(shù)據(jù)庫進(jìn)行查詢
在查詢階段,我們使用Milvus來存儲(chǔ)GraphRAG本地搜索中實(shí)體描述的向量嵌入。
這種方法將知識圖譜的結(jié)構(gòu)化數(shù)據(jù)與輸入文檔的非結(jié)構(gòu)化數(shù)據(jù)相結(jié)合,為LLM提供了額外的相關(guān)實(shí)體信息,從而能夠得出更準(zhǔn)確的答案。
import os
import pandas as pd
import tiktoken
from graphrag.query.context_builder.entity_extraction import EntityVectorStoreKey
from graphrag.query.indexer_adapters import (
# read_indexer_covariates,
read_indexer_entities,
read_indexer_relationships,
read_indexer_reports,
read_indexer_text_units,
)
from graphrag.query.input.loaders.dfs import (
store_entity_semantic_embeddings,
)
from graphrag.query.llm.oai.chat_openai import ChatOpenAI
from graphrag.query.llm.oai.embedding import OpenAIEmbedding
from graphrag.query.llm.oai.typing import OpenaiApiType
from graphrag.query.question_gen.local_gen import LocalQuestionGen
from graphrag.query.structured_search.local_search.mixed_context import (
LocalSearchMixedContext,
)
from graphrag.query.structured_search.local_search.search import LocalSearch
from graphrag.vector_stores import MilvusVectorStore
output_dir = os.path.join(index_root, "output")
subdirs = [os.path.join(output_dir, d) for d in os.listdir(output_dir)]
latest_subdir = max(subdirs, key=os.path.getmtime) # 獲取最新的輸出目錄
INPUT_DIR = os.path.join(latest_subdir, "artifacts")
COMMUNITY_REPORT_TABLE = "create_final_community_reports"
ENTITY_TABLE = "create_final_nodes"
ENTITY_EMBEDDING_TABLE = "create_final_entities"
RELATIONSHIP_TABLE = "create_final_relationships"
COVARIATE_TABLE = "create_final_covariates"
TEXT_UNIT_TABLE = "create_final_text_units"
COMMUNITY_LEVEL = 2
7、從索引過程中加載數(shù)據(jù)
在索引過程中,會(huì)生成幾個(gè)parquet文件。我們將其加載到內(nèi)存中,并將實(shí)體描述信息存儲(chǔ)在Milvus向量數(shù)據(jù)庫中。
讀取實(shí)體:
# 讀取節(jié)點(diǎn)表以獲取社區(qū)和度量數(shù)據(jù)
entity_df = pd.read_parquet(f"{INPUT_DIR}/{ENTITY_TABLE}.parquet")
entity_embedding_df = pd.read_parquet(f"{INPUT_DIR}/{ENTITY_EMBEDDING_TABLE}.parquet")
entities = read_indexer_entities(entity_df, entity_embedding_df, COMMUNITY_LEVEL)
description_embedding_store = MilvusVectorStore(
collection_name="entity_description_embeddings",
)
# description_embedding_store.connect(uri="http://localhost:19530") # 用于Milvus docker服務(wù)
description_embedding_store.connect(uri="./milvus.db") # For Milvus Lite
entity_description_embeddings = store_entity_semantic_embeddings(
entities=entities, vectorstore=description_embedding_store
)
print(f"實(shí)體數(shù)量:{len(entity_df)}")
entity_df.head()
實(shí)體數(shù)量:651
圖片
讀取關(guān)系
relationship_df = pd.read_parquet(f"{INPUT_DIR}/{RELATIONSHIP_TABLE}.parquet")
relationships = read_indexer_relationships(relationship_df)
print(f"關(guān)系數(shù)量: {len(relationship_df)}")
relationship_df.head()
關(guān)系數(shù)量:290
讀取社區(qū)報(bào)告
report_df = pd.read_parquet(f"{INPUT_DIR}/{COMMUNITY_REPORT_TABLE}.parquet")
reports = read_indexer_reports(report_df, entity_df, COMMUNITY_LEVEL)
print(f"報(bào)告記錄:{len(report_df)}")
report_df.head()
報(bào)告記錄:45
讀取文本單元
text_unit_df = pd.read_parquet(f"{INPUT_DIR}/{TEXT_UNIT_TABLE}.parquet")
text_units = read_indexer_text_units(text_unit_df)
print(f"文本單元記錄:{len(text_unit_df)}")
text_unit_df.head()
文本單元記錄:51
8、構(gòu)建本地搜索引擎
一切準(zhǔn)備就緒,現(xiàn)在已經(jīng)有了本地搜索引擎所需的所有數(shù)據(jù)。接下來,我們可以利用這些數(shù)據(jù),配合一個(gè)大型語言模型(LLM)和一個(gè)嵌入模型,來構(gòu)建一個(gè)??LocalSearch?
?實(shí)例,為搜索任務(wù)提供強(qiáng)大的支持。
api_key = os.environ["OPENAI_API_KEY"] # 你的OpenAI API密鑰
llm_model = "gpt-4o" # 或 gpt-4-turbo-preview
embedding_model = "text-embedding-3-small"
llm = ChatOpenAI(
api_key=api_key,
model=llm_model,
api_type=OpenaiApiType.OpenAI,
max_retries=20,
)
token_encoder = tiktoken.get_encoding("cl100k_base")
text_embedder = OpenAIEmbedding(
api_key=api_key,
api_base=None,
api_type=OpenaiApiType.OpenAI,
model=embedding_model,
deployment_name=embedding_model,
max_retries=20,
)
context_builder = LocalSearchMixedContext(
community_reports=reports,
text_units=text_units,
entities=entities,
relatinotallow=relationships,
covariates=None, #covariates,#todo
entity_text_embeddings=description_embedding_store,
embedding_vectorstore_key=EntityVectorStoreKey.ID, # 如果向量存儲(chǔ)使用實(shí)體標(biāo)題作為ID,則將此設(shè)置為EntityVectorStoreKey.TITLE
text_embedder=text_embedder,
token_encoder=token_encoder,
)
local_context_params = {
"text_unit_prop": 0.5,
"community_prop": 0.1,
"conversation_history_max_turns": 5,
"conversation_history_user_turns_only": True,
"top_k_mapped_entities": 10,
"top_k_relationships": 10,
"include_entity_rank": True,
"include_relationship_weight": True,
"include_community_rank": False,
"return_candidate_context": False,
"embedding_vectorstore_key": EntityVectorStoreKey.ID, # 如果向量存儲(chǔ)使用實(shí)體標(biāo)題作為ID,則將此設(shè)置為EntityVectorStoreKey.TITLE
"max_tokens": 12
_000, # 根據(jù)你的模型的令牌限制更改此設(shè)置(如果你使用的是8k限制的模型,一個(gè)好設(shè)置可能是5000)
}
llm_params = {
"max_tokens": 2_000, # 根據(jù)你的模型的令牌限制更改此設(shè)置(如果你使用的是8k限制的模型,一個(gè)好設(shè)置可能是1000=1500)
"temperature": 0.0,
}
search_engine = LocalSearch(
llm=llm,
context_builder=context_builder,
token_encoder=token_encoder,
llm_params=llm_params,
context_builder_params=local_context_params,
response_type="multiple paragraphs", # 描述響應(yīng)類型和格式的自由形式文本,可以是任何內(nèi)容,例如優(yōu)先列表、單段、多段、多頁報(bào)告
)
9、進(jìn)行查詢
result = await search_engine.asearch("Tell me about Leonardo Da Vinci")
print(result.response)
# 萊昂納多·達(dá)·芬奇
萊昂納多·達(dá)·芬奇,1452年出生于佛羅倫薩附近的文奇鎮(zhèn),被廣泛譽(yù)為意大利文藝復(fù)興時(shí)期最多才多藝的天才之一。他的全名是萊昂納多·迪·塞爾·皮耶羅·達(dá)·安東尼奧·迪·塞爾·皮耶羅·迪·塞爾·圭多·達(dá)·芬奇,他是塞爾·皮耶羅的非婚生和長子,塞爾·皮耶羅是一位鄉(xiāng)村公證人[數(shù)據(jù):實(shí)體(0]。萊昂納多的貢獻(xiàn)涵蓋了藝術(shù)、科學(xué)、工程和哲學(xué)等多個(gè)領(lǐng)域,他被譽(yù)為基督教時(shí)代最萬能的天才[數(shù)據(jù):實(shí)體(8]。
## 早期生活和訓(xùn)練
萊昂納多早期的才華得到了他父親的認(rèn)可,他父親將他的一些畫作帶給了安德烈亞·德爾·維羅基奧,一位著名的藝術(shù)家和雕塑家。維羅基奧對萊昂納多的才華印象深刻,于1469-1470年左右接受了他進(jìn)入自己的工作室。在這里,萊昂納多遇到了包括博蒂切利和洛倫佐·迪·克雷迪在內(nèi)的其他著名藝術(shù)家[數(shù)據(jù):來源(6, 7]。到1472年,萊昂納多被佛羅倫薩畫家行會(huì)錄取,標(biāo)志著他職業(yè)生涯的開始[數(shù)據(jù):來源(7)]。
## 藝術(shù)杰作
萊昂納多或許以他的標(biāo)志性畫作最為人所知,如《蒙娜麗莎》和《最后的晚餐》?!睹赡塞惿芬云湮⒚畹谋砬楹驮敿?xì)的背景而聞名,現(xiàn)藏于盧浮宮,仍然是世界上最著名的藝術(shù)品之一[數(shù)據(jù):關(guān)系(0, 45]。《最后的晚餐》是一幅壁畫,描繪了耶穌宣布他的一個(gè)門徒將背叛他的那一刻,位于米蘭圣瑪利亞·格拉齊教堂的餐廳[數(shù)據(jù):來源(2]。其他重要作品包括《巖間圣母》和他大約在1489-1490年開始的《繪畫論》[數(shù)據(jù):關(guān)系(7, 12]。
## 科學(xué)和工程貢獻(xiàn)
萊昂納多的天才超越了藝術(shù),延伸到各種科學(xué)和工程事業(yè)。他在解剖學(xué)、光學(xué)和水力學(xué)方面做出了重要觀察,他的筆記本里充滿了預(yù)示許多現(xiàn)代發(fā)明的草圖和想法。例如,他預(yù)示了哥白尼關(guān)于地球運(yùn)動(dòng)的理論和拉馬克對動(dòng)物的分類[數(shù)據(jù):關(guān)系(38, 39]。他對光影法則和明暗對比的掌握對藝術(shù)和科學(xué)都產(chǎn)生了深遠(yuǎn)影響[數(shù)據(jù):來源(45)]。
## 贊助和職業(yè)關(guān)系
萊昂納多的職業(yè)生涯受到他的贊助人的重大影響。米蘭公爵盧多維科·斯福爾扎雇傭萊昂納多作為宮廷畫家和普通工匠,委托了各種作品,甚至在1499年贈(zèng)送給他一個(gè)葡萄園[數(shù)據(jù):關(guān)系(9, 19, 84]。在他的晚年,萊昂納多在法國國王弗朗西斯一世的贊助下搬到了法國,國王為他提供了豐厚的收入,并對他評價(jià)很高[數(shù)據(jù):關(guān)系(114, 37]。萊昂納多在法國安博瓦茲附近的克洛克斯莊園度過了他最后幾年,國王經(jīng)常拜訪他,他得到了他的密友和助手弗朗切斯科·梅爾齊的支持[數(shù)據(jù):關(guān)系(28, 122)]。
## 遺產(chǎn)和影響
萊昂納多·達(dá)·芬奇的影響遠(yuǎn)遠(yuǎn)超出了他的一生。他在米蘭創(chuàng)立了一所繪畫學(xué)校,他的技術(shù)和教導(dǎo)被他的學(xué)生和追隨者,如喬瓦尼·安布羅焦·達(dá)·普雷迪斯和弗朗切斯科·梅爾齊傳承下去[數(shù)據(jù):關(guān)系(6, 15, 28]。他的作品繼續(xù)受到慶祝和研究,鞏固了他作為文藝復(fù)興時(shí)期最偉大的大師之一的遺產(chǎn)。萊昂納多將藝術(shù)和科學(xué)融合的能力在兩個(gè)領(lǐng)域都留下了不可磨滅的印記,激勵(lì)著無數(shù)的藝術(shù)家和科學(xué)家[數(shù)據(jù):實(shí)體(148, 86]; 關(guān)系(27, 12)]。
總之,萊昂納多·達(dá)·芬奇對藝術(shù)、科學(xué)和工程的無與倫比的貢獻(xiàn),加上他的創(chuàng)新思維和對同時(shí)代人及后代的深遠(yuǎn)影響,使他成為人類成就史上的一位杰出人物。他的遺產(chǎn)繼續(xù)激發(fā)著欽佩和研究,強(qiáng)調(diào)了他的天才的永恒相關(guān)性。
GraphRAG的結(jié)果具體,明確標(biāo)出了引用的數(shù)據(jù)源。
10、問題生成
GraphRAG還可以根據(jù)歷史查詢生成問題,這對于在聊天機(jī)器人對話中創(chuàng)建推薦問題非常有用。這種方法結(jié)合了知識圖譜的結(jié)構(gòu)化數(shù)據(jù)和輸入文檔的非結(jié)構(gòu)化數(shù)據(jù),產(chǎn)生與特定實(shí)體相關(guān)的候選問題。
question_generator = LocalQuestionGen(
llm=llm,
context_builder=context_builder,
token_encoder=token_encoder,
llm_params=llm_params,
context_builder_params=local_context_params,
)
question_history = [
"Tell me about Leonardo Da Vinci",
"Leonardo's early works",
]
Generate questions based on history.
candidate_questions = await question_generator.agenerate(
question_history=question_history, context_data=None, question_count=5
)
candidate_questions.response
["- 萊昂納多·達(dá)·芬奇的早期作品有哪些,它們存放在哪里?",
"萊昂納多·達(dá)·芬奇與安德烈亞·德爾·維羅基奧的關(guān)系如何影響了他的早期作品?",
"萊昂納多·達(dá)·芬奇在米蘭期間承擔(dān)了哪些重要項(xiàng)目?",
"萊昂納多·達(dá)·芬奇的工程技能如何促成他的項(xiàng)目?",
"萊昂納多·達(dá)·芬奇與法國弗朗西斯一世的關(guān)系有何重要性?"]
如果你想刪除索引以節(jié)省空間,可以移除索引根。
# import shutil
#
# shutil.rmtree(index_root)
11、結(jié)語
本文帶領(lǐng)大家深入了解了GraphRAG技術(shù),這是一種融合知識圖譜來強(qiáng)化RAG應(yīng)用的創(chuàng)新手段。GraphRAG特別擅長處理那些需要跨信息片段進(jìn)行多步驟推理和全面回答問題的復(fù)雜任務(wù)。
結(jié)合Milvus向量數(shù)據(jù)庫后,GraphRAG能夠高效地在龐大的數(shù)據(jù)集中探索復(fù)雜的語義聯(lián)系,從而得出更精準(zhǔn)、更深刻的分析結(jié)果。這種強(qiáng)強(qiáng)聯(lián)合的解決方案,使GraphRAG成為眾多實(shí)際通用人工智能(GenAI)應(yīng)用中的得力助手,為理解和處理復(fù)雜信息提供了強(qiáng)有力的支持。
本文轉(zhuǎn)載自公眾號AI科技論談
原文鏈接:??https://mp.weixin.qq.com/s/rfTwh0B6wnKjauiZsTqwyw??
