使用本地HuggingFace模型+LangGraph时如何仅获取模型输出？

问题原因 #

你使用的TinyLlama-1.1B-Chat-v1.0模型依赖特定对话格式标记（<|system|>、<|user|>、<|assistant|>），当前代码未正确配置模型生成参数，也未严格遵循模型要求的对话模板，导致模型返回包含完整上下文的内容，而非仅助手回复。

解决方案 #

1. 核心修改方向 #

• 手动加载模型对应的Tokenizer，获取正确的终止/填充标记ID
• 配置生成参数，禁止返回完整输入文本
• 使用符合模型格式的对话模板，明确角色分隔

修改后的完整代码 #

from typing import Annotated
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langchain_core.messages import HumanMessage
from langchain_huggingface import ChatHuggingFace
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

# 加载本地模型与Tokenizer
model_id = "../downloaded_models/TinyLlama-1.1B-Chat-v1.0"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)

# 配置生成参数：仅返回新生成内容，设置终止标记
generation_config = {
    "max_new_tokens": 512,
    "eos_token_id": tokenizer.eos_token_id,
    "pad_token_id": tokenizer.pad_token_id,
    "return_full_text": False,
    "temperature": 0.7
}

# 创建文本生成Pipeline
hf_pipeline = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    generation_config=generation_config
)

llm = ChatHuggingFace(llm=hf_pipeline)

# 构建符合TinyLlama格式的对话模板
prompt_template = ChatPromptTemplate.from_messages([
    ("system", "<|system|>\n你是一个乐于助人的助手</s>"),
    ("user", "<|user|>\n{input}</s>"),
    ("assistant", "<|assistant|>")
])

# 绑定模板与LLM
chain = prompt_template | llm

class State(TypedDict):
    messages: Annotated[list, add_messages]

graph_builder = StateGraph(State)

def chatbot(state: State):
    user_msg = state["messages"][-1].content
    ai_msg = chain.invoke({"input": user_msg})
    return {"messages": [ai_msg]}

graph_builder.add_node("chatbot", chatbot)
graph_builder.add_edge(START, "chatbot")
graph_builder.add_edge("chatbot", END)
graph = graph_builder.compile()

while True:
    user_input = input("User: ")
    if user_input.lower() in ["quit", "exit", "q"]:
        print("再见！")
        break
    for event in graph.stream({"messages": HumanMessage(content=user_input)}):
        for value in event.values():
            print("Assistant:", value["messages"][-1].content)

关键修改说明 #

• Tokenizer加载：确保获取模型原生的标记ID，避免生成内容混乱
• 生成参数配置：return_full_text=False 让模型仅返回新增的回复内容，不包含输入的prompt部分
• 对话模板：严格遵循TinyLlama的角色标记格式，引导模型专注生成助手回复
• Chain绑定：将模板与LLM绑定，自动处理每次输入的格式转换，无需手动拼接标记

效果验证 #

修改后输入问题，模型输出将仅包含助手的回复内容，不再附带用户输入和角色标识。

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