现在,让我们转向一个实际应用场景——流式输出。
在我的LLM项目开发过程中,第一个遇到的需求就是实现流式输出。这个功能可以减少用户的等待时间,从而提升用户体验。不过,目前作为一个演示项目,GraphRAG还没有支持流式输出。
因此,我将带领大家实现这个流式输出功能。我们的目标不仅是添加新功能,更重要的是通过这个过程,加深你对GraphRAG源码的理解。
命令行添加streaming参数
GraphRAG在graphrag/query/main.py文件中包含了执行Query的命令行入口代码,为了提升它的功能并且让其可以进行流式输出,我们计划增加一个名为"streaming"的参数。
这个新增的参数可以通过下面的代码段添加:
parser.add_argument( "--streaming", help="Whether to output the response in a streaming format", action="store_true", )
此代码段的作用是在命令行解析器(parser)中添加一个新的选项"–streaming"。当这个选项被设置时(即在命令行中加上–streaming),则该参数的值为True。
我们希望这样的更新可以改变程序的处理方式,使其在执行查询时能够以流格式逐步输出结果,而不是等待所有操作完成后一次性输出。
当我们所有的代码写完之后,执行下面的命令可以轻松地实现流式输出:
poetry run poe query --root ./ragtest --streaming --method global '路飞有哪些伙伴?'
run_local_search & run_global_search增加streaming参数
在GraphRAG源码中Local Query和Global Query调用的函数分别是graphrag/query/cli.py里的run_local_search和run_global_search。现在我们需要为这两个函数添加入参:
def run_local_search( config_dir: str | None, data_dir: str | None, root_dir: str | None, community_level: int, response_type: str, streaming: bool, query: str, ): ....
def run_global_search( config_dir: str | None, data_dir: str | None, root_dir: str | None, community_level: int, response_type: str, streaming: bool, query: str, ): ...
紧接着我们需要在graphrag/query/main.py的__main__
中把streaming这个命令行参数传给这两个函数:
... match args.method: case SearchType.LOCAL: run_local_search( args.config, args.data, args.root, args.community_level, args.response_type, args.streaming, args.query[0], ) case SearchType.GLOBAL: run_global_search( args.config, args.data, args.root, args.community_level, args.response_type, args.streaming, args.query[0], ) case _: raise ValueError(INVALID_METHOD_ERROR)
在GraphRAG的框架中,run_local_search和run_global_search函数的执行逻辑主要包括以下几个步骤:
-
根据构建的知识图谱各种parquet文件,组织执行search函数所需要的数据。
-
利用这些数据,构建一个执行查询的对象,我们将它称为search_engine。
-
调用search_engine对象的search或asearch方法来执行实际的查询操作。
在这三个步骤中,第一步和第二步负责数据的准备和查询对象的构建,并不会受到我们是否引入streaming参数的影响。换句话说,无论用户是否选择了流式输出,知识图谱数据的组织和search_engine对象的构建都是相同的。
因此,我们真正需要修改的部分是第三步:调用search_engine对象的search或asearch方法:
result = await search_engine.asearch(query=query) reporter.success(f"Local Search Response: {result.response}") return result.response
如果用户选择了流式输出,那么我们就需要修改这两个方法,使其能够返回一个可迭代对象,从而可以实现逐步推送查询结果的功能。同时,也要保证当用户没有选择流式输出时,这两个方法仍然能够按照原来的方式直接返回查询结果。
为此,我新增一个astream_search方法用于处理流式输出。
在接下来的讨论中,我会以local_search
函数为例来详细说明我们的流式功能,对global_search
函数的优化也将遵循相同的方法,因为global_search
和local_search
函数本质上是非常相似的。它们之间唯一的区别在于所使用的搜索引擎对象不同:global_search
使用的是GlobalSearch
引擎,而local_search
使用的是LocalSearch
引擎。
本次修改的GraphRAG的版本是v0.3.0
local_search
在GraphRAG的最新代码版本中,run_local_search
函数调用了位于graphrag/query/api.py文件中的local_search
方法。此方法封装了我们之前所提的第二步(构建查询引擎对象)和第三步(执行查询操作)的逻辑。
为了实现流式输出功能,我们需要对这个local_search
方法进行一些修改。具体的改动将包括在用户选择了流式输出时使其返回一个可迭代对象,从而实现逐步推送查询结果的功能。同时,我们也需要保证当用户没有选择流式输出时,local_search
方法仍然能像原来那样直接返回查询结果。
search_engine = get_local_search_engine( config=config, reports=read_indexer_reports(community_reports, nodes, community_level), text_units=read_indexer_text_units(text_units), entities=_entities, relationships=read_indexer_relationships(relationships), covariates={"claims": _covariates}, description_embedding_store=description_embedding_store, response_type=response_type, ) if not streaming: result = await search_engine.asearch(query=query) reporter.success(f"Global Search Response: {result.response}") return result.response else: import sys full_resp = '' results = search_engine.astream_search(query=query) reporter.success(f'Global Search Response: \n') async for result in results: sys.stdout.write(result) sys.stdout.flush() full_resp += result sys.stdout.write('\b\n') return full_resp
astream_search
为了实现流式响应功能,我们在代码中新增了一个名为astream_search
的方法。这个方法与原有的asearch
方法有着相似的功能,但它们之间的主要区别在于如何调用LLM。
当使用asearch
方法时,所有的查询结果会在一次性完成后返回给用户。换句话说,LLM会处理完所有的查询请求后才将结果返回。但当我们使用astream_search
方法时,LLM将会在处理每一部分查询请求后就立即返回当前的结果。也就是说,astream_search
调用了LLM的流式响应功能。
async def astream_search( self, query: str, conversation_history: ConversationHistory | None = None, **kwargs, ) -> SearchResult: """Build local search context that fits a single context window and generate answer for the user query.""" start_time = time.time() search_prompt = "" context_text, context_records = self.context_builder.build_context( query=query, conversation_history=conversation_history, **kwargs, **self.context_builder_params, ) log.info("GENERATE ANSWER: %s. QUERY: %s", start_time, query) try: search_prompt = self.system_prompt.format( context_data=context_text, response_type=self.response_type ) search_messages = [ {"role": "system", "content": search_prompt}, {"role": "user", "content": query}, ] response = await self.llm.astream_generate( messages=search_messages, callbacks=self.callbacks, **self.llm_params, ) return SearchResult( response=response, context_data=context_records, context_text=context_text, completion_time=time.time() - start_time, llm_calls=1, prompt_tokens=num_tokens(search_prompt, self.token_encoder), ) except Exception: log.exception("Exception in _asearch") return SearchResult( response="", context_data=context_records, context_text=context_text, completion_time=time.time() - start_time, llm_calls=1, prompt_tokens=num_tokens(search_prompt, self.token_encoder), )
astream_generate
为了实现流式查询,我们需要在代码的多个部分进行修改。具体来说,在源码中search_engine.asearch
实际上是调用了self.llm.agenerate
方法。因此,为了支持流式查询,我们也需要在相应的位置添加一个名为astream_generate
的新方法。
这个新方法将被添加到位于graphrag/query/llm/oai/chat_openai.py文件中的ChatOpenAI类中。添加astream_generate
的目的是为了让ChatOpenAI类能够生成流式响应。
async def astream_generate( self, messages: str | list[Any], callbacks: list[BaseLLMCallback] | None = None, **kwargs: Any, ) -> AsyncGenerator[str, None] | None: """Generate text asynchronously with streaming.""" try: retryer = AsyncRetrying( stop=stop_after_attempt(self.max_retries), wait=wait_exponential_jitter(max=10), reraise=True, retry=retry_if_exception_type(self.retry_error_types), # type: ignore ) async for attempt in retryer: with attempt: return self._astream_generate( messages=messages, callbacks=callbacks, **kwargs, ) except RetryError as e: self._reporter.error(f"Error at astream_generate(): {e}") return else: return async def _astream_generate( self, messages: str | list[Any], callbacks: list[BaseLLMCallback] | None = None, **kwargs: Any, ) -> AsyncGenerator[str, None]: model = self.model if not model: raise ValueError(_MODEL_REQUIRED_MSG) response = await self.async_client.chat.completions.create( # type: ignore model=model, messages=messages, # type: ignore stream=True, **kwargs, ) async for chunk in response: if not chunk or not chunk.choices: continue delta = ( chunk.choices[0].delta.content if chunk.choices[0].delta and chunk.choices[0].delta.content else "" ) # type: ignore yield delta if callbacks: for callback in callbacks: callback.on_llm_new_token(delta)
测试
经过上面代码的新增和修改之后,下面我们就可以来测试streaming功能了,运行很流畅,pefect
`poetry run poe query --root ./ragtest --streaming --method local '路飞有哪些伙伴?'`
总结
我们已经成功地实现了streaming功能,并且在这个过程中,只添加了为数不多的一些新的代码。然而,实现这个功能不是我们真正想要达到的目标,更重要的是通过这个过程深入理解和掌握GraphRAG的源码。
如何学习大模型 AI ?
由于新岗位的生产效率,要优于被取代岗位的生产效率,所以实际上整个社会的生产效率是提升的。
但是具体到个人,只能说是:
“最先掌握AI的人,将会比较晚掌握AI的人有竞争优势”。
这句话,放在计算机、互联网、移动互联网的开局时期,都是一样的道理。
我在一线互联网企业工作十余年里,指导过不少同行后辈。帮助很多人得到了学习和成长。
我意识到有很多经验和知识值得分享给大家,也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑,所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限,很多互联网行业朋友无法获得正确的资料得到学习提升,故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。
第一阶段(10天):初阶应用
该阶段让大家对大模型 AI有一个最前沿的认识,对大模型 AI 的理解超过 95% 的人,可以在相关讨论时发表高级、不跟风、又接地气的见解,别人只会和 AI 聊天,而你能调教 AI,并能用代码将大模型和业务衔接。
- 大模型 AI 能干什么?
- 大模型是怎样获得「智能」的?
- 用好 AI 的核心心法
- 大模型应用业务架构
- 大模型应用技术架构
- 代码示例:向 GPT-3.5 灌入新知识
- 提示工程的意义和核心思想
- Prompt 典型构成
- 指令调优方法论
- 思维链和思维树
- Prompt 攻击和防范
- …
第二阶段(30天):高阶应用
该阶段我们正式进入大模型 AI 进阶实战学习,学会构造私有知识库,扩展 AI 的能力。快速开发一个完整的基于 agent 对话机器人。掌握功能最强的大模型开发框架,抓住最新的技术进展,适合 Python 和 JavaScript 程序员。
- 为什么要做 RAG
- 搭建一个简单的 ChatPDF
- 检索的基础概念
- 什么是向量表示(Embeddings)
- 向量数据库与向量检索
- 基于向量检索的 RAG
- 搭建 RAG 系统的扩展知识
- 混合检索与 RAG-Fusion 简介
- 向量模型本地部署
- …
第三阶段(30天):模型训练
恭喜你,如果学到这里,你基本可以找到一份大模型 AI相关的工作,自己也能训练 GPT 了!通过微调,训练自己的垂直大模型,能独立训练开源多模态大模型,掌握更多技术方案。
到此为止,大概2个月的时间。你已经成为了一名“AI小子”。那么你还想往下探索吗?
- 为什么要做 RAG
- 什么是模型
- 什么是模型训练
- 求解器 & 损失函数简介
- 小实验2:手写一个简单的神经网络并训练它
- 什么是训练/预训练/微调/轻量化微调
- Transformer结构简介
- 轻量化微调
- 实验数据集的构建
- …
第四阶段(20天):商业闭环
对全球大模型从性能、吞吐量、成本等方面有一定的认知,可以在云端和本地等多种环境下部署大模型,找到适合自己的项目/创业方向,做一名被 AI 武装的产品经理。
- 硬件选型
- 带你了解全球大模型
- 使用国产大模型服务
- 搭建 OpenAI 代理
- 热身:基于阿里云 PAI 部署 Stable Diffusion
- 在本地计算机运行大模型
- 大模型的私有化部署
- 基于 vLLM 部署大模型
- 案例:如何优雅地在阿里云私有部署开源大模型
- 部署一套开源 LLM 项目
- 内容安全
- 互联网信息服务算法备案
- …
学习是一个过程,只要学习就会有挑战。天道酬勤,你越努力,就会成为越优秀的自己。
如果你能在15天内完成所有的任务,那你堪称天才。然而,如果你能完成 60-70% 的内容,你就已经开始具备成为一名大模型 AI 的正确特征了。