NLTK vs LLM：GraphRAG实体提取方法深度大比

2024年9月4日修改

作者：还就可 | 深入LLM Agent应用开发

原文： https://mp.weixin.qq.com/s/OjGv-7p4...

书接上文《 75% 成本和时间削减：优化微软 GraphRAG 索引的秘密》，我们在修复一些bug和缩减大量成本后，使得GraphRAG更易于使用。其实在缩减成本上，GraphaRAG还有一个隐藏的功能，那就是使用NLTK来提取实体。NLTK，全称Natural Language Toolkit（自然语言工具包），是一个开源的Python库，主要用于自然语言处理（NLP）的研究和开发。这篇文章主要讨论了如何通过使用NLTK和LLM两种不同的方法来进行实体提取，以优化和比较在GraphRAG中的应用效果。

1. GraphRAG配置NLTK

NLTK由Steven Bird和Edward Loper等人在宾夕法尼亚大学计算机和信息科学系开发，它为NLP研究者和开发者提供了一系列的工具、数据集和教程。NLTK包含的功能非常广泛，涵盖了自然语言处理的多个方面，包括但不限于：​

•
分词（Tokenization）：将文本分割成单词或短语。​

•
词性标注（Part-of-Speech Tagging）：为每个词标注其语法类别。​

•
命名实体识别（Named Entity Recognition）：识别文本中的专有名词，如人名、地名、组织名等。​

•
句法分析（Syntactic Parsing）：分析句子的结构和语法关系。​

•
语义分析（Semantic Analysis）：理解文本的深层含义。​

•
情感分析（Sentiment Analysis）：识别文本中表达的情感倾向。​

•
文本分类（Text Classification）：自动分类文本到预定义的类别。​

•
机器翻译（Machine Translation）：将文本从一种语言翻译成另一种语言。​

•
问答系统（Question Answering）：自动回答文本提出的问题。​

由于其全面性和易用性，NLTK成为了教育和研究领域中最受欢迎的NLP库之一，尤其适合初学者学习自然语言处理的基础。在GraphRAG中，只需配置settings.yaml即可，找到 entity_extraction 按照如下配置：​

代码块

entity_extraction:​
  strategy:​
    type: "nltk"​
  ## llm: override the global llm settings for this task​
  ## parallelization: override the global parallelization settings for this task​
  ## async_mode: override the global async_mode settings for this task​
  prompt: "prompts/entity_extraction.txt"​
  entity_types: [organization,person,geo,event]​
  max_gleanings: 0​

如果你使用的是OpenAI的LLM服务，这样配置即可。但如果你使用的是其他LLM服务，这个配置是无法生效的，这里还存一个bug，我已经提交了 PR [1] 进行修复，应该很快就会合并。实体提取的代码在目录 graphrag/index/verbs/entities/extraction/strategies/graph_intelligence ，默认策略是 graph_intelligence ，也就是使用LLM根据Prompt进行自动提取。上文有说到实体提取的Prompt很大（位于根目录 prompts/entity_extraction.txt ），大约有2000多个Token，并且和文档的chunk的数量成倍数关系。因此，如果采用NLTK进行分词，将对成本进一步优化，岂不妙哉？​

以下测试都采用Google的Gemma2-9b小模型，主要是groq的免费额度较大，而且它跟随指令比较好，基本上每次都能完整跑完索引。​

2. 测试NLTK实体提取

以英国大文豪狄更斯的小说《圣诞颂歌》为例，将其存储为txt文件，放入input文件夹。​

代码块

curl https://www.gutenberg.org/cache/epub/24022/pg24022.txt > ./input/book.txt

开始索引

代码块

poetry run poe index --root .

由于篇幅相对较长，大约46155个Tokens，耗时2149秒完成，其中实体提取耗时936秒，社区报告1200秒。顺便提一句，整个索引过程所耗时间，可以在输出文件夹 output/{日期}/artfacts/stats.json 文件中看到，他统计了每个workflow的时间以及总共耗时。本次索引输出文件都在 output/20240716-104546 ，方便下次指定使用。​

代码块

{​
    "total_runtime": 2149.660534143448,​
    "num_documents": 1,​
    "input_load_time": 0,​
    "workflows": {​
        "create_base_text_units": {​
            "overall": 0.25486016273498535,​
            "0_orderby": 0.0034270286560058594,​
            "1_zip": 0.0015499591827392578,​
            "2_aggregate_override": 0.003968715667724609,​
            "3_chunk": 0.13766002655029297,​
            "4_select": 0.002666950225830078,​
            "5_unroll": 0.0044748783111572266,​
            "6_rename": 0.002095937728881836,​
            "7_genid": 0.008743047714233398,​
            "8_unzip": 0.003268003463745117,​
            "9_copy": 0.0034067630767822266,​
            "10_filter": 0.08233785629272461​
        },​
   ...​
    }​
}​

NLTK vs LLM：GraphRAG实体提取方法深度大比​

NLTK vs LLM：GraphRAG实体提取方法深度大比