[Python] Llama 2 with Langchain tools

EDIT:

I found that it works with Llama 2 70b, but not with Llama 2 13b. Llama 2 13b uses the tool correctly and observes the final answer which is in its agent_scratchpad, but it outputs an empty string at the end whereas Llama 2 70b outputs 'It looks like the answer is 18.37917367995256!' which is correct.

Original post:

I am trying to follow this tutorial on using Llama 2 with Langchain tools (you don't have to look at the tutorial all code is contained in this question).

My code is very similar to that in the tutorial except I am using a local model rather than connecting to Hugging Face and I am not using bitsandbytes for quantisation since it requires cuda and I am on macOS.

I am using the unquantised Meta Llama 2 13b chat model meta-llama/Llama-2-13b-chat-hf.

The model appears to be outputting JSON correctly but for some reason I am getting "Could not parse LLM output".

Here is the code (added \ before triple backtick due to Stackoverflow code formatting). Note I added the following to the prompt "When Assistant responds with JSON they make sure to enclose the JSON with three back ticks." This resolves an issue where the model was outputting JSON that was not in the correct format (surrounded with backticks and a 'json' tag).

import transformers
model_id = './Models/Llama_2/llama_2_13b'

# initialize the model
model = transformers.AutoModelForCausalLM.from_pretrained(model_id)
tokenizer = transformers.AutoTokenizer.from_pretrained(model_id)

generate_text = transformers.pipeline(
model=model, tokenizer=tokenizer,
return_full_text=True, # langchain expects the full text
task='text-generation',
# we pass model parameters here too
temperature=0.01, # 'randomness' of outputs, 0.0 is the min and 1.0 the max
max_new_tokens=512, # mex number of tokens to generate in the output
repetition_penalty=1.1 # without this output begins repeating
)

from langchain.llms import HuggingFacePipeline

llm = HuggingFacePipeline(pipeline=generate_text)

from langchain.memory import ConversationBufferWindowMemory
from langchain.agents import load_tools

memory = ConversationBufferWindowMemory(
memory_key="chat_history", k=5, return_messages=True, output_key="output"
)

tools = load_tools(["llm-math"], llm=llm)

from langchain.agents import initialize_agent

# initialize agent
agent = initialize_agent(
agent="chat-conversational-react-description",
tools=tools,
llm=llm,
verbose=True,
early_stopping_method="generate",
memory=memory,
handle_parsing_errors=True
)

# special tokens used by llama 2 chat
B_INST, E_INST = "[INST]", "[/INST]"
B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"

# create the system message
sys_msg = "<s>" + B_SYS + """Assistant is a expert JSON builder designed to assist with a wide range of tasks.

Assistant is able to respond to the User and use tools using JSON strings that contain "action" and "action_input" parameters.

All of Assistant's communication is performed using this JSON format.

Assistant can also use tools by responding to the user with tool use instructions in the same "action" and "action_input" JSON format. Tools available to Assistant are:

- "Calculator": Useful for when you need to answer questions about math.
- To use the calculator tool, Assistant should write like so:
```json
{{"action": "Calculator",
"action_input": "sqrt(4)"}}
```

When Assistant responds with JSON they make sure to enclose the JSON with three back ticks.

Here are some previous conversations between the Assistant and User:

User: Hey how are you today?
Assistant: ```json
{{"action": "Final Answer",
"action_input": "I'm good thanks, how are you?"}}
\```
User: I'm great, what is the square root of 4?
Assistant: ```json
{{"action": "Calculator",
"action_input": "sqrt(4)"}}
\```
User: 2.0
Assistant: ```json
{{"action": "Final Answer",
"action_input": "It looks like the answer is 2!"}}
\```
User: Thanks could you tell me what 4 to the power of 2 is?
Assistant: ```json
{{"action": "Calculator",
"action_input": "4**2"}}
\```
User: 16.0
Assistant: ```json
{{"action": "Final Answer",
"action_input": "It looks like the answer is 16!"}}
\```

Here is the latest conversation between Assistant and User.""" + E_SYS

new_prompt = agent.agent.create_prompt(
system_message=sys_msg,
tools=tools
)

agent.agent.llm_chain.prompt = new_prompt

instruction = B_INST + " Respond to the following in JSON with 'action' and 'action_input' values " + E_INST
human_msg = instruction + "\nUser: {input}"

agent.agent.llm_chain.prompt.messages[2].prompt.template = human_msg



Then I prompt it with;

agent("hey how are you today?")


and I get the following;

> Entering new AgentExecutor chain...


Assistant: ```json
{"action": "Final Answer",
"action_input": "I'm good thanks, how are you?"}
\```

> Finished chain.
{'input': 'hey how are you today?',
'chat_history': [],
'output': "I'm good thanks, how are you?"}


Which is great, but when I prompt it with a question that requires use of a tool;

agent("what is 4 to the power of 2.1?")


I get;

> Entering new AgentExecutor chain...


Assistant: ```json
{"action": "Calculator",
"action_input": "4**2.1"}
\```
Observation: Answer: 18.37917367995256
Thought:Could not parse LLM output:
Observation: Invalid or incomplete response
Thought:Could not parse LLM output:
Observation: Invalid or incomplete response
Thought:Could not parse LLM output:
Observation: Invalid or incomplete response


and it just gets stuck in a loop repeating "Could not parse LLM output" and "Invalid or incomplete response"

Does anyone know how to fix the "Could not parse LLM output" errors?

Supposedly this code worked for the authors of the tutorial.

I am using Python 3.11.5 with Anaconda, tensorflow 2.15.0, transformers 4.35.2, langchain 0.0.336, on macOS Sonoma.

EDIT:

I updated the code to use the output parser from here.

from langchain.agents import AgentOutputParser
from langchain.agents.conversational_chat.prompt import FORMAT_INSTRUCTIONS
from langchain.output_parsers.json import parse_json_markdown
from langchain.schema import AgentAction, AgentFinish

class OutputParser(AgentOutputParser):
def get_format_instructions(self) -> str:
return FORMAT_INSTRUCTIONS

def parse(self, text: str) -> AgentAction | AgentFinish:
try:
# this will work IF the text is a valid JSON with action and action_input
response = parse_json_markdown(text)
action, action_input = response["action"], response["action_input"]
if action == "Final Answer":
# this means the agent is finished so we call AgentFinish
return AgentFinish({"output": action_input}, text)
else:
# otherwise the agent wants to use an action, so we call AgentAction
return AgentAction(action, action_input, text)
except Exception:
# sometimes the agent will return a string that is not a valid JSON
# often this happens when the agent is finished
# so we just return the text as the output
return AgentFinish({"output": text}, text)

@property
def _type(self) -> str:
return "conversational_chat"

# initialize output parser for agent
parser = OutputParser()


by adding

agent_kwargs={"output_parser": parser}


to initialise_agent, and the code no longer produces an error but the final output from the LLM is still empty.

> Entering new AgentExecutor chain...

Assistant: ```json
{"action": "Calculator",
"action_input": "4**2.1"}
\```
Observation: Answer: 18.37917367995256
Thought:

> Finished chain.
{'input': 'what is 4 to the power of 2.1?',
'chat_history': [HumanMessage(content='hey how are you today?'),
AIMessage(content="I'm good thanks, how are you?")],
'output': ''}

Continue reading...