Phase 1: 文档预处理 (论文 3.1.1)

在 Nano GraphRAG 里，预处理文档的任务在 GraphRAG.ainsert() 中完成。GraphRAG.ainsert() 由 GraphRAG.insert() 调用，并且使用了 asyncio.get_event_loop() 保证执行完毕。

文档去重

ainsert() 首先计算文档的 MD5 哈希值，检查文档是否已经添加过了。如果已经添加过了，那么就不用再插入数据库。

new_docs 首先准备所有待检查的字符串，计算其哈希值。

1
2
3
4
new_docs = {
    compute_mdhash_id(c.strip(), prefix="doc-"): {"content": c.strip()}
    for c in string_or_strings
}

然后丢入 self.full_docs 检查是否有重复的字符串并剔除

1
_add_doc_keys = await self.full_docs.filter_keys(list(new_docs.keys()))

剔除完成后的字符串拼接成 { hash: str } 形式的字典重新赋值回 new_docs.

1
new_docs = {k: v for k, v in new_docs.items() if k in _add_doc_keys}

如果没有新文档，直接退出。然后输出一点调试日志

1
2
3
4
if not len(new_docs):
    logger.warning(f"All docs are already in the storage")
    return
logger.info(f"[New Docs] inserting {len(new_docs)} docs")

文档分块 Text Chunk

否则，准备插入文档。先对文档 get_chunk() 拆分成若干个 text chunks

1
2
3
4
5
6
7
8
9
10
11
12
13
inserting_chunks = get_chunks(
    new_docs=new_docs,
    # 要切分的文档  

    chunk_func=self.chunk_func,
    # chunk 切分方法，默认按 token 数量

    overlap_token_size=self.chunk_overlap_token_size,
    # chunk 之间重叠的 token 数量，充当上下文的作用

    max_token_size=self.chunk_token_size,
    # 一个 chunk 最多多少 token
)

然后 get_chunk() 加载 Tokenizer (这里调用的是 OpenAI tiktoken 的库)，将 docs（纯文本）转化为 Tokens，再调用 chunk_func 进行切块，每一块都标注成

1
2
3
4
5
6
{
    "tokens": token 个数,
    "content": token 内容（已经不是文本了）, 
    "chunk_order_index": 是第几个 chunk,
    "full_doc_id": 所属文档的哈希值,
}

标注完成后，再为每一个 text chunk 计算哈希值 chunk-xxxxx，放入 inserting_chunk 并返回

Text Chunk 去重

对 text chunks 也进行去重

1
2
3
4
5
6
7
8
9
10
11
12
13
_add_chunk_keys = await self.text_chunks.filter_keys(
    list(inserting_chunks.keys())
)
inserting_chunks = {
    k: v for k, v in inserting_chunks.items() if k in _add_chunk_keys
}
if not len(inserting_chunks):
    logger.warning(f"All chunks are already in the storage")
    return
logger.info(f"[New Chunks] inserting {len(inserting_chunks)} chunks")
if self.enable_naive_rag:
    logger.info("Insert chunks for naive RAG")
    await self.chunks_vdb.upsert(inserting_chunks)

和文档去重的逻辑很相似，就是在 KVStorage 里按哈希值查找，去重

更新数据库

这里有另一点比较重要的是，如果需要插入新的 chunk，那么首先需要把 self.community_reports 清空。因为插入新块后，可能这个文档的 community 就会改变

1
await self.community_reports.drop()
References