Langchain结合本地Transformer小模型与Agent/Tools解析异常咨询

问题核心是DistilGPT2这类轻量Transformer模型无法严格遵循LangChain默认Agent的结构化输出格式要求，导致解析器抛出OutputParserException。以下是针对该问题的具体解决方案：

1. 自定义简化输出解析器 #

默认Agent解析器要求输出严格遵循Thought/Action/Action Input/Observation的嵌套结构，小模型很难精准生成。我们可以自定义一个容错性更高的解析器，用正则匹配提取关键信息：

from langchain.agents import AgentOutputParser
from langchain.schema import AgentAction, AgentFinish
import re

class SimpleDistilGPT2Parser(AgentOutputParser):
    def parse(self, llm_output: str) -> AgentAction | AgentFinish:
        # 匹配Final Answer的情况
        if "Final Answer:" in llm_output:
            return AgentFinish(
                return_values={"output": llm_output.split("Final Answer:")[-1].strip()},
                log=llm_output,
            )
        # 匹配Action和Action Input的情况（兼容小模型的松散输出）
        action_match = re.search(r"Action:\s*(.*?)\s*Action Input:\s*(.*)", llm_output, re.DOTALL)
        if action_match:
            action = action_match.group(1).strip()
            action_input = action_match.group(2).strip().strip('"')
            return AgentAction(tool=action, tool_input=action_input, log=llm_output)
        # 如果匹配失败，直接返回Final Answer（容错处理）
        return AgentFinish(
            return_values={"output": llm_output.strip()},
            log=llm_output,
        )

2. 优化Prompt模板，强化格式引导 #

针对小模型的理解能力，简化Prompt中的格式说明，用更直白的指令重复格式要求：

from langchain.prompts import StringPromptTemplate
from langchain.chains import LLMChain

class SimpleDistilGPT2Prompt(StringPromptTemplate):
    template: str
    tools: list

    def format(self, **kwargs) -> str:
        # 获取中间步骤的历史
        intermediate_steps = kwargs.pop("intermediate_steps")
        thoughts = ""
        for action, observation in intermediate_steps:
            thoughts += f"Action: {action.tool}\nAction Input: {action.tool_input}\nObservation: {observation}\n"
        kwargs["thoughts"] = thoughts
        # 生成工具说明（简化格式）
        tool_strings = "\n".join([f"- {tool.name}: {tool.description}" for tool in self.tools])
        kwargs["tools"] = tool_strings
        kwargs["tool_names"] = ", ".join([tool.name for tool in self.tools])
        return self.template.format(**kwargs)

# 定义简化后的Prompt内容
PROMPT_TEMPLATE = """完成以下任务，你可以使用提供的工具。
可用工具：
{tools}

格式要求（必须严格遵守）：
如果需要使用工具，输出：
Action: [工具名称]
Action Input: [工具的输入内容]

如果已经得到答案，输出：
Final Answer: [最终答案]

当前任务：{input}
历史步骤：{thoughts}
你的回答："""

3. 初始化适配小模型的Agent #

将自定义的解析器和Prompt整合到Agent中，替换默认组件：

from langchain.agents import AgentExecutor, LLMSingleActionAgent
from langchain.llms import HuggingFacePipeline
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

# 加载DistilGPT2模型和管道
model_name = "distilgpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=100,
    temperature=0.1,  # 降低温度提升输出稳定性
    top_p=0.95,
    repetition_penalty=1.15
)
llm = HuggingFacePipeline(pipeline=pipe)

# 初始化Prompt和解析器
your_tools_list = []  # 替换成你的工具列表
prompt = SimpleDistilGPT2Prompt(
    template=PROMPT_TEMPLATE,
    tools=your_tools_list,
    input_variables=["input", "intermediate_steps"]
)
output_parser = SimpleDistilGPT2Parser()

# 创建Agent并执行
agent = LLMSingleActionAgent(
    llm_chain=LLMChain(llm=llm, prompt=prompt),
    output_parser=output_parser,
    stop=["\nObservation:"],  # 定义停止词，避免模型生成冗余内容
    allowed_tools=[tool.name for tool in your_tools_list]
)
agent_executor = AgentExecutor.from_agent_and_tools(agent=agent, tools=your_tools_list, verbose=True)

# 测试执行
agent_executor.run("你的任务内容")

关键注意事项 #

• 降低小模型的生成温度（temperature），减少输出随机性，提升格式一致性。
• 限制最大生成长度（max_new_tokens），避免模型生成过多无关内容。
• 加入重复惩罚（repetition_penalty），防止模型重复输出无意义内容。

内容的提问来源于stack exchange，提问作者Dipesh Paul