上文我们简单介绍了一下fast-graphrag，今天我们初步跑一下demo以及看看背后的存储结果。
数据准备
验证阶段倒也不需要直接下载官方的book.txt来跑，毕竟100多k的文本，处理要点token。
这边我直接用模型随便生成了几个故事。
book.txt
在遥远的阿瓦隆大陆，有一个被遗忘的古老森林，那里居住着神秘的生物和拥有魔力的植物。一天，年轻的冒险者艾莉娅在森林的边缘发现了一条隐藏的小径，这条小径引导她来到了一个被迷雾笼罩的湖泊。湖中央有一座孤岛，岛上耸立着一座古老的塔楼。艾莉娅划着小船靠近岛屿，决心揭开塔楼的秘密。当她踏入塔楼的那一刻，她的命运与古老的预言交织在一起，她必须找到传说中的光之水晶，以阻止即将到来的黑暗势力。
在2150年，地球已经不再是人类唯一的家园。人类已经在火星和几个小行星上建立了殖民地。主角杰克是一位星际货运船的船长，他的船“星际穿梭者”号在一次例行运输任务中遭遇了未知的宇宙现象。这次遭遇导致他的船只被迫降落在一个未知的星球上。这个星球表面覆盖着奇异的水晶结构，似乎是一种外星文明的遗迹。杰克和他的船员必须解开这些结构的秘密，同时寻找回家的路。
在繁华的都市中，有一家名为“时光咖啡馆”的小店，它的老板是一位温文尔雅的中年人，名叫林峰。这家咖啡馆有一个不为人知的秘密：它的地下室是一个时间旅行者的聚集地。一天，一位神秘的女子出现在咖啡馆，她声称自己是从未来来的旅行者，需要林峰的帮助来修复她的时间机器。林峰被卷入了一个跨越时空的冒险，他必须帮助这位女子修复机器，同时保护咖啡馆的秘密不被外界发现。
在一座古老的庄园中，一场盛大的化妆舞会正在进行。然而，当午夜钟声响起时，庄园的主人被发现死在了自己的书房中。所有的嫌疑人都穿着华丽的服装，戴着面具，没有人知道谁才是真正的凶手。侦探艾米丽接手了这个案件，她必须通过观察和推理来揭开每个人的秘密，找出真正的凶手。随着调查的深入，庄园中隐藏的秘密逐渐浮出水面，而真相往往比表面看起来更加复杂。
测试代码
官方demo上面有一步是export openai的环境变量，这个其实目前看起来是行不通的。真要跑起来还是要使用自定义的llm模型的配置。如下：在创建GraphRAG的时候传入自定义的llm_service 和 embedding_service
from fast_graphrag import GraphRAG
import os
from fast_graphrag._llm import OpenAIEmbeddingService, OpenAILLMService
import openai
import instructor
DOMAIN = "Analyze this story and identify the characters. Focus on how they interact with each other, the locations they explore, and their relationships."
ENTITY_TYPES = ["Character", "Animal", "Place", "Object", "Activity", "Event"]
grag = GraphRAG(
working_dir="./book_example",
domain=DOMAIN,
example_queries="",
entity_types=ENTITY_TYPES,
n_checkpoints=2,
config=GraphRAG.Config(
llm_service=OpenAILLMService(
model="gpt-4o-mini", base_url=os.getenv("OPENAI_API_BASE"), api_key=os.getenv("OPENAI_API_KEY"), mode=instructor.Mode.JSON
),
embedding_service=OpenAIEmbeddingService(
model="text-embedding-ada-002",
base_url=os.getenv("OPENAI_API_BASE"),
api_key=os.getenv("OPENAI_API_KEY"),
embedding_dim=1536, # the output embedding dim of the chosen model
),
),
)
with open("./book.txt") as f:
data = f.read()
print(data)
grag.insert(data)
# print(grag.query("艾丽西亚是谁?他在干啥").response)
跑完之后会在本地book_example目录下面生成1个文件夹（因为我设置了n_checkpoints=2）和6个文件。
entities_hnsw_metadata.pkl
map_r2c_blob_data.pkl
chunks_kv_data.pkl # 分片文本数据
graph_igraph_data.pklz # 存储图结构（节点和边）
entities_hnsw_index_1536.bin
map_e2r_blob_data.pkl
# HNSW（Hierarchical Navigable Small World）索引，用于高效的近似最近邻搜索，通常在推荐系统或相似性搜索中使用
pkl文件和pklz文件都可以通过pickle包来打开。
import pickle
import os
import gzip
# 替换为您的 .pkl 文件路径
file_path = "./book_example"
filenames = os.listdir(file_path)
for filename in filenames:
if filename.endswith('.pkl'):
print(filename)
with open(f"{file_path}/{filename}", 'rb') as file:
data = pickle.load(file)
print(data)
print('==='*10)
elif filename.endswith('.pklz'): # 处理 .pklz 文件
print(filename)
with gzip.open(f"{file_path}/{filename}", 'rb') as file:
data = pickle.load(file) # 使用 pklz 加载数据
print(data)
print('==='*10)
可以看到图结构那一块，因为我生成的四个故事是一次性给他insert的，导致他把整个这块作为了整体，四个故事的主角也关联起来了。
模型提取出的实体包括
...
{
"name": "林峰",
"type": "Character",
"desc": "时光咖啡馆的老板"
}
...
查询
grag.query("咖啡馆老板是谁").response
# 咖啡馆老板是林峰。
第一步，他根据用户的query，获取到了如下信息：
{"named": ["[PERSON] 咖啡馆老板"], "generic": ["咖啡馆"]}
根据这个信息去图检索、向量搜索拼凑出了如下提示词：
You are a helpful assistant analyzing the given input data to provide an helpful response to the user query.
# INPUT DATA
## Entities
```csv
name description
林峰 时光咖啡馆的老板
神秘女子 声称是未来来的旅行者
时光咖啡馆 一家咖啡馆，地下室是时间旅行者的聚集地
《此处我省略了很多》
```
## Relationships
```csv
source target description
林峰 神秘女子 林峰帮助神秘女子修复时间机器
林峰 时光咖啡馆 林峰是时光咖啡馆的老板
时光咖啡馆 神秘女子 神秘女子出现在时光咖啡馆寻求帮助
《此处我省略了很多》
```
## Sources
[1] 在遥远的阿瓦隆大陆，有一个被遗忘的古老森林，那里居住着神秘的生物和拥有魔力的植物。一天，年轻的冒险者艾莉娅在森林的边缘发现了一条隐藏的小径，这条小径引导她来到了一个被迷雾笼罩的湖泊。湖中央有一座孤岛，岛上耸立着一座古老的塔楼。艾莉娅划着小船靠近岛屿，决心揭开塔楼的秘密。当她踏入塔楼的那一刻，她的命运与古老的预言交织在一起，她必须找到传说中的光之水晶，以阻止即将到来的黑暗势力。在2150年，地球已经不再是人类唯一的家园。人类已经在火星和几个小行星上建立了殖民地。主角杰克是一位星际货运船的船长，他的船“星际穿梭者”号在一次例行运输任务中遭遇了未知的宇宙现象。这次遭遇导致他的船只被迫降落在一个未知的星球上。这个星球表面覆盖着奇异的水晶结构，似乎是一种外星文明的遗迹。杰克和他的船员必须解开这些结构的秘密，同时寻找回家的路。在繁华的都市中，有一家名为“时光咖啡馆”的小店，它的老板是一位温文尔雅的中年人，名叫林峰。这家咖啡馆有一个不为人知的秘密：它的地下室是一个时间旅行者的聚集地。一天，一位神秘的女子出现在咖啡馆，她声称自己是从未来来的旅行者，需要林峰的帮助来修复她的时间机器。林峰被卷入了一个跨越时空的冒险，他必须帮助这位女子修复机器，同时保护咖啡馆的秘密不被外界发现。在一座古老的庄园中，一场盛大的化妆舞会正在进行。然而，当午夜钟声响起时，庄园的主人被发现死在了自己的书房中。所有的嫌疑人都穿着华丽的服装，戴着面具，没有人知道谁才是真正的凶手。侦探艾米丽接手了这个案件，她必须通过观察和推理来揭开每个人的秘密，找出真正的凶手。随着调查的深入，庄园中隐藏的秘密逐渐浮出水面，而真相往往比表面看起来更加复杂。
=====
# USER QUERY
咖啡馆老板是谁
# INSTRUCTIONS
Your goal is to provide a response to the user query using the relevant information in the input data:
- the "Entities" and "Relationships" tables contain high-level information. Use these tables to identify the most important entities and relationships to respond to the query.
- the "Sources" list contains raw text sources to help answer the query. It may contain noisy data, so pay attention when analyzing it.
Follow these steps:
1. Read and understand the user query.
2. Look at the "Entities" and "Relationships" tables to get a general sense of the data and understand which information is the most relevant to answer the query.
3. Carefully analyze all the "Sources" to get more detailed information. Information could be scattered across several sources, use the identified relevant entities and relationships to guide yourself through the analysis of the sources.
4. Write the response to the user query based on the information you have gathered. Be very concise and answer the user query directly. If the response cannot be inferred from the input data, just say no relevant information was found. Do not make anything up or add unrelevant information.
Answer:
最后得到了答案。
拓展
使用 .pkl 文件的原因主要有以下几点：
序列化和反序列化：.pkl 文件是 Python 的 pickle 模块生成的文件格式，用于将 Python 对象序列化为字节流，方便存储和传输。反之，可以将字节流反序列化为原始对象。
高效存储：与其他文本格式（如 CSV 或 JSON）相比，.pkl 文件可以更高效地存储复杂的数据结构，如嵌套列表、字典和自定义对象。
保留数据类型：使用 pickle 可以保留数据的类型信息，确保在加载数据时能够恢复为原始的 Python 对象，而不是简单的字符串或数字。
方便使用：对于需要频繁读写的机器学习模型或数据集，使用 .pkl 文件可以简化代码，减少数据处理的复杂性。
到这，本篇已经完结，咱们后面再慢慢研究这个fast graphrag，看看源代码，也欢迎大家在评论区留言，分享你对模型应用的一些经验和看法~
