Llama2Chat Repeats Input Instead of Generating Relevant Responses for the Last <INST> Block

[JustAname12345]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

from langgraph.graph import MessagesState, StateGraph, START
from langgraph.graph.message import BaseMessage, add_messages
from langchain.schema import HumanMessage, AIMessage, SystemMessage
from transformers import AutoTokenizer, AutoModelForCausalLM
from langgraph.checkpoint.memory import MemorySaver
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from typing_extensions import TypedDict, Annotated
from typing import Sequence
import torch
from langchain_core.messages import SystemMessage, trim_messages

class State(TypedDict):
    messages: Annotated[Sequence[BaseMessage], add_messages]
    language: str

prompt_template = ChatPromptTemplate.from_messages([
    (
        'system',
        " You are a knowledgeable assistant. Answer all questions in {language}. Keep your responses concise and natural."
    ),
    MessagesPlaceholder(variable_name="messages")
])

model_name = "/root/autodl-tmp/llama2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")

def token_counter(messages):
    input_str = "".join(
        [
            f"{message.content}"
            for message in messages
            if isinstance(message, (SystemMessage, HumanMessage, AIMessage))
        ]
    )
    return len(tokenizer.encode(input_str))

trimmer = trim_messages(
    max_tokens=150,
    strategy="last",
    token_counter=token_counter,
    include_system=True,
    allow_partial=False,
    start_on="human",
)


def format_prompt(messages, system_prompt):
    
    formatted = []
    first_message = True
    for i in range(len(messages)):
        message = messages[i]

        if isinstance(message, HumanMessage):
            if first_message:
                formatted.append(
                    f"<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n{message.content} [/INST]"
                )
                first_message = False
            else:
                formatted.append(f"<s>[INST] {message.content} [/INST]")

        elif isinstance(message, AIMessage):
            formatted[-1] += f" {message.content} </s>"

    if not formatted[-1].endswith("</s>") and isinstance(messages[-1], HumanMessage):
        formatted[-1] += ""

    return "\n".join(formatted)

    
async def call_model(state: State):
    trimmed_messages = trimmer.invoke(state["messages"])
   
    print("===== Trimmed Messages =====")
    for msg in trimmed_messages:
        print(f"{msg.__class__.__name__}: {msg.content}")
        
    state["messages"] = trimmed_messages  
    input_prompt = prompt_template.invoke(
        {"messages": trimmed_messages, "language": state["language"]}
    )
    system_message = ""
    input_prompt = format_prompt(trimmed_messages, system_message)
    print("===== Input Prompt =====")
    print(input_prompt)
    print("=======================")
    
    inputs = tokenizer(input_prompt, return_tensors="pt", truncation=True, max_length=128).to("cuda")
    outputs = model.generate(
        **inputs,
        max_new_tokens=128,
        num_beams=5,
        temperature=0.7,
        
    )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    #input(response)
    ai_message = response.strip()
    input(ai_message)
    return {"messages": messages + [AIMessage(content=ai_message)]}

workflow = StateGraph(state_schema=State)
workflow.add_edge(START, "model")
workflow.add_node("model", call_model)
memory = MemorySaver()
app = workflow.compile(checkpointer=memory)


messages = [
    HumanMessage(content="hi! I'm bob"),
    AIMessage(content="hi!"),
    HumanMessage(content="I like vanilla ice cream"),
    AIMessage(content="nice"),
    HumanMessage(content="what does does x,x,x,y mean?"),
    AIMessage(content="means nothing"),
    HumanMessage(content="thanks"),
    AIMessage(content="no problem!"),
    HumanMessage(content="having fun?"),
    AIMessage(content="yes!"),
]

async def main():
    config = {"configurable": {"thread_id": "abc123"}}
    query1 = "what is my name?"
    state = {
        "messages": messages,
        "language": 'english'
    }
    state["messages"].extend([HumanMessage(content=query1)])
    output = await app.ainvoke(state, config)
    print("\nResponse to first question:")
    output['messages'][-1].pretty_print()
    
    query2 = "what does x,x,x,y mean?"
    state["messages"] = trimmer.invoke(output["messages"])  # 只保留最近的几轮对话
    state["messages"].extend([HumanMessage(content=query2)])
    output = await app.ainvoke(state, config)
    print("\nResponse to second question:")
    output["messages"][-1].pretty_print()

if __name__ == "__main__":
    import asyncio
    asyncio.run(main())

Description

when I utilize a conversational state to interact with a Hugging Face Llama2 model. The goal is to answer the latest question from a sequence of inputs. However, the model often repeats the input text instead of generating a meaningful response. like following

System Info

System Information

OS: Linux
OS Version: #96-Ubuntu SMP Wed Sep 20 08:23:49 UTC 2023
Python Version: 3.12.8 | packaged by Anaconda, Inc. | (main, Dec 11 2024, 16:31:09) [GCC 11.2.0]

Package Information

langchain_core: 0.3.29
langchain: 0.3.14
langsmith: 0.2.10
langchain_text_splitters: 0.3.5
langgraph_sdk: 0.1.51

Optional packages not installed

langserve

Other Dependencies

aiohttp: 3.11.11
async-timeout: Installed. No version info available.
httpx: 0.28.1
jsonpatch: 1.33
langsmith-pyo3: Installed. No version info available.
numpy: 2.2.1
orjson: 3.10.14
packaging: 24.2
pydantic: 2.10.5
PyYAML: 6.0.2
requests: 2.32.3
requests-toolbelt: 1.0.0
SQLAlchemy: 2.0.37
tenacity: 9.0.0
typing-extensions: 4.12.2
zstandard: Installed. No version info available.