{
"cells": [
{
"cell_type": "markdown",
"id": "ce4a9ccf-4bd6-43fb-a24d-b6a7da401a96",
"metadata": {},
"source": [
"## Load xLAM model"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b1351d81-4502-4b65-b88a-464acd0e80f8",
"metadata": {},
"outputs": [],
"source": [
"import torch \n",
"from transformers import AutoModelForCausalLM, AutoTokenizer\n",
"torch.random.manual_seed(0) \n",
"\n",
"model_name = \"Salesforce/xLAM-7b-r\"\n",
"model = AutoModelForCausalLM.from_pretrained(model_name, device_map=\"auto\", torch_dtype=\"auto\", trust_remote_code=True)\n",
"tokenizer = AutoTokenizer.from_pretrained(model_name) "
]
},
{
"cell_type": "markdown",
"id": "2cdd5bae-da43-4713-9956-360f1f3a9721",
"metadata": {},
"source": [
"## Build the prompt"
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "e138e9f6-0543-427c-bce6-b4f14765a040",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"import json\n",
"\n",
"# Please use our provided instruction prompt for best performance\n",
"task_instruction = \"\"\"\n",
"Based on the previous context and API request history, generate an API request or a response as an AI assistant.\"\"\".strip()\n",
"\n",
"format_instruction = \"\"\"\n",
"The output should be of the JSON format, which specifies a list of generated function calls. The example format is as follows, please make sure the parameter type is correct. If no function call is needed, please make \n",
"tool_calls an empty list \"[]\".\n",
"```\n",
"{\"thought\": \"the thought process, or an empty string\", \"tool_calls\": [{\"name\": \"api_name1\", \"arguments\": {\"argument1\": \"value1\", \"argument2\": \"value2\"}}]}\n",
"```\n",
"\"\"\".strip()\n",
"\n",
"get_weather_api = {\n",
" \"name\": \"get_weather\",\n",
" \"description\": \"Get the current weather for a location\",\n",
" \"parameters\": {\n",
" \"type\": \"object\",\n",
" \"properties\": {\n",
" \"location\": {\n",
" \"type\": \"string\",\n",
" \"description\": \"The city and state, e.g. San Francisco, New York\"\n",
" },\n",
" \"unit\": {\n",
" \"type\": \"string\",\n",
" \"enum\": [\"celsius\", \"fahrenheit\"],\n",
" \"description\": \"The unit of temperature to return\"\n",
" }\n",
" },\n",
" \"required\": [\"location\"]\n",
" }\n",
"}\n",
"\n",
"search_api = {\n",
" \"name\": \"search\",\n",
" \"description\": \"Search for information on the internet\",\n",
" \"parameters\": {\n",
" \"type\": \"object\",\n",
" \"properties\": {\n",
" \"query\": {\n",
" \"type\": \"string\",\n",
" \"description\": \"The search query, e.g. 'latest news on AI'\"\n",
" }\n",
" },\n",
" \"required\": [\"query\"]\n",
" }\n",
"}\n",
"\n",
"openai_format_tools = [get_weather_api, search_api]\n",
"\n",
"# Define the input query and available tools\n",
"query = \"What's the weather like in New York in fahrenheit?\"\n",
"\n",
"# Helper function to convert openai format tools to our more concise xLAM format\n",
"def convert_to_xlam_tool(tools):\n",
" ''''''\n",
" if isinstance(tools, dict):\n",
" return {\n",
" \"name\": tools[\"name\"],\n",
" \"description\": tools[\"description\"],\n",
" \"parameters\": {k: v for k, v in tools[\"parameters\"].get(\"properties\", {}).items()}\n",
" }\n",
" elif isinstance(tools, list):\n",
" return [convert_to_xlam_tool(tool) for tool in tools]\n",
" else:\n",
" return tools\n",
"\n",
"def build_conversation_history_prompt(conversation_history: str):\n",
" parsed_history = []\n",
" for step_data in conversation_history:\n",
" parsed_history.append({\n",
" \"step_id\": step_data[\"step_id\"],\n",
" \"thought\": step_data[\"thought\"],\n",
" \"tool_calls\": step_data[\"tool_calls\"],\n",
" \"next_observation\": step_data[\"next_observation\"],\n",
" \"user_input\": step_data['user_input']\n",
" })\n",
" \n",
" history_string = json.dumps(parsed_history)\n",
" return f\"\\n[BEGIN OF HISTORY STEPS]\\n{history_string}\\n[END OF HISTORY STEPS]\\n\"\n",
" \n",
" \n",
"# Helper function to build the input prompt for our model\n",
"def build_prompt(task_instruction: str, format_instruction: str, tools: list, query: str, conversation_history: list):\n",
" prompt = f\"[BEGIN OF TASK INSTRUCTION]\\n{task_instruction}\\n[END OF TASK INSTRUCTION]\\n\\n\"\n",
" prompt += f\"[BEGIN OF AVAILABLE TOOLS]\\n{json.dumps(xlam_format_tools)}\\n[END OF AVAILABLE TOOLS]\\n\\n\"\n",
" prompt += f\"[BEGIN OF FORMAT INSTRUCTION]\\n{format_instruction}\\n[END OF FORMAT INSTRUCTION]\\n\\n\"\n",
" prompt += f\"[BEGIN OF QUERY]\\n{query}\\n[END OF QUERY]\\n\\n\"\n",
" \n",
" if len(conversation_history) > 0: prompt += build_conversation_history_prompt(conversation_history)\n",
" return prompt\n",
"\n",
"\n",
" \n",
"# Build the input and start the inference\n",
"xlam_format_tools = convert_to_xlam_tool(openai_format_tools)\n",
"\n",
"conversation_history = []\n",
"content = build_prompt(task_instruction, format_instruction, xlam_format_tools, query, conversation_history)\n",
"\n",
"messages=[\n",
" { 'role': 'user', 'content': content}\n",
"]\n"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "ff7bccd5-fa04-4fbe-92b3-13f58914da4d",
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[BEGIN OF TASK INSTRUCTION]\n",
"Based on the previous context and API request history, generate an API request or a response as an AI assistant.\n",
"[END OF TASK INSTRUCTION]\n",
"\n",
"[BEGIN OF AVAILABLE TOOLS]\n",
"[{\"name\": \"get_weather\", \"description\": \"Get the current weather for a location\", \"parameters\": {\"location\": {\"type\": \"string\", \"description\": \"The city and state, e.g. San Francisco, New York\"}, \"unit\": {\"type\": \"string\", \"enum\": [\"celsius\", \"fahrenheit\"], \"description\": \"The unit of temperature to return\"}}}, {\"name\": \"search\", \"description\": \"Search for information on the internet\", \"parameters\": {\"query\": {\"type\": \"string\", \"description\": \"The search query, e.g. 'latest news on AI'\"}}}]\n",
"[END OF AVAILABLE TOOLS]\n",
"\n",
"[BEGIN OF FORMAT INSTRUCTION]\n",
"The output should be of the JSON format, which specifies a list of generated function calls. The example format is as follows, please make sure the parameter type is correct. If no function call is needed, please make \n",
"tool_calls an empty list \"[]\".\n",
"```\n",
"{\"thought\": \"the thought process, or an empty string\", \"tool_calls\": [{\"name\": \"api_name1\", \"arguments\": {\"argument1\": \"value1\", \"argument2\": \"value2\"}}]}\n",
"```\n",
"[END OF FORMAT INSTRUCTION]\n",
"\n",
"[BEGIN OF QUERY]\n",
"What's the weather like in New York in fahrenheit?\n",
"[END OF QUERY]\n",
"\n",
"\n"
]
}
],
"source": [
"print(content)"
]
},
{
"cell_type": "markdown",
"id": "a5fb0006-9f5d-4d79-a8cd-819bad627441",
"metadata": {},
"source": [
"## Get the model output (agent_action)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "cbe56588-c786-4913-9062-373a22a92e08",
"metadata": {},
"outputs": [],
"source": [
"inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors=\"pt\").to(model.device)\n",
"\n",
"# tokenizer.eos_token_id is the id of <|EOT|> token\n",
"outputs = model.generate(inputs, max_new_tokens=512, do_sample=False, num_return_sequences=1, eos_token_id=tokenizer.eos_token_id)\n",
"agent_action = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)\n"
]
},
{
"cell_type": "markdown",
"id": "b20ed2ae-86f6-489b-ad54-fe7ea911667b",
"metadata": {},
"source": [
"For demo purpose, we use an example agent_action"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "ab20c084-44fa-403d-92a5-1b8ced72e9be",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"agent_action = \"\"\"{\"thought\": \"\", \"tool_calls\": [{\"name\": \"get_weather\", \"arguments\": {\"location\": \"New York\"}}]}\n",
"\"\"\".strip()"
]
},
{
"cell_type": "markdown",
"id": "1cd4d8e4-ee6b-499e-b75f-a48df7848a60",
"metadata": {},
"source": [
"### Add follow-up question"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "825649ba-2691-43a2-b3d8-7baf8b66d46e",
"metadata": {},
"outputs": [],
"source": [
"def parse_agent_action(agent_action: str):\n",
" \"\"\"\n",
" Given an agent's action, parse it to add to conversation history\n",
" \"\"\"\n",
" try: parsed_agent_action_json = json.loads(agent_action)\n",
" except: return \"\", []\n",
" \n",
" if \"thought\" not in parsed_agent_action_json.keys(): thought = \"\"\n",
" else: thought = parsed_agent_action_json[\"thought\"]\n",
" \n",
" if \"tool_calls\" not in parsed_agent_action_json.keys(): tool_calls = []\n",
" else: tool_calls = parsed_agent_action_json[\"tool_calls\"]\n",
" \n",
" return thought, tool_calls\n",
"\n",
"def update_conversation_history(conversation_history: list, agent_action: str, environment_response: str, user_input: str):\n",
" \"\"\"\n",
" Update the conversation history list based on the new agent_action, environment_response, and/or user_input\n",
" \"\"\"\n",
" thought, tool_calls = parse_agent_action(agent_action)\n",
" new_step_data = {\n",
" \"step_id\": len(conversation_history) + 1,\n",
" \"thought\": thought,\n",
" \"tool_calls\": tool_calls,\n",
" \"next_observation\": environment_response,\n",
" \"user_input\": user_input,\n",
" }\n",
" \n",
" conversation_history.append(new_step_data)\n",
"\n",
"def get_environment_response(agent_action: str):\n",
" \"\"\"\n",
" Get the environment response for the agent_action\n",
" \"\"\"\n",
" # TODO: add custom implementation here\n",
" error_message, response_message = \"\", \"Sunny, 81 degrees\"\n",
" return {\"error\": error_message, \"response\": response_message}\n",
"\n"
]
},
{
"cell_type": "markdown",
"id": "051e6aff-c21b-4dcb-9eb8-c34154d90c39",
"metadata": {},
"source": [
"1. **Get the next state after agent's response:**\n",
" The next 2 lines are examples of getting environment response and user_input.\n",
" It is depended on particular usage, we can have either one or both of those."
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "649a8e9d-9757-408c-9214-0590556c2db4",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"environment_response = get_environment_response(agent_action)\n",
"user_input = \"Now, search on the Internet for cute puppies\""
]
},
{
"cell_type": "markdown",
"id": "9c9c9418-1c54-4381-81d1-7f3834037739",
"metadata": {},
"source": [
"2. After we got environment_response and (or) user_input, we want to add to our conversation history"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "bcfe89f3-8237-41bf-b92c-7c7568366042",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"[{'step_id': 1,\n",
" 'thought': '',\n",
" 'tool_calls': [{'name': 'get_weather',\n",
" 'arguments': {'location': 'New York'}}],\n",
" 'next_observation': {'error': '', 'response': 'Sunny, 81 degrees'},\n",
" 'user_input': 'Now, search on the Internet for cute puppies'}]"
]
},
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"update_conversation_history(conversation_history, agent_action, environment_response, user_input)\n",
"conversation_history"
]
},
{
"cell_type": "markdown",
"id": "23ba97c6-2356-49e8-a07b-0e664b7f505c",
"metadata": {},
"source": [
"3. We now can build the prompt with the updated history, and prepare the inputs for the LLM"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "ed204b3a-3be5-431b-b355-facaf31309d2",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"content = build_prompt(task_instruction, format_instruction, xlam_format_tools, query, conversation_history)\n",
"messages=[\n",
" { 'role': 'user', 'content': content}\n",
"]\n"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "8af843aa-6a47-4938-a455-567ea0cccce3",
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[BEGIN OF TASK INSTRUCTION]\n",
"Based on the previous context and API request history, generate an API request or a response as an AI assistant.\n",
"[END OF TASK INSTRUCTION]\n",
"\n",
"[BEGIN OF AVAILABLE TOOLS]\n",
"[{\"name\": \"get_weather\", \"description\": \"Get the current weather for a location\", \"parameters\": {\"location\": {\"type\": \"string\", \"description\": \"The city and state, e.g. San Francisco, New York\"}, \"unit\": {\"type\": \"string\", \"enum\": [\"celsius\", \"fahrenheit\"], \"description\": \"The unit of temperature to return\"}}}, {\"name\": \"search\", \"description\": \"Search for information on the internet\", \"parameters\": {\"query\": {\"type\": \"string\", \"description\": \"The search query, e.g. 'latest news on AI'\"}}}]\n",
"[END OF AVAILABLE TOOLS]\n",
"\n",
"[BEGIN OF FORMAT INSTRUCTION]\n",
"The output should be of the JSON format, which specifies a list of generated function calls. The example format is as follows, please make sure the parameter type is correct. If no function call is needed, please make \n",
"tool_calls an empty list \"[]\".\n",
"```\n",
"{\"thought\": \"the thought process, or an empty string\", \"tool_calls\": [{\"name\": \"api_name1\", \"arguments\": {\"argument1\": \"value1\", \"argument2\": \"value2\"}}]}\n",
"```\n",
"[END OF FORMAT INSTRUCTION]\n",
"\n",
"[BEGIN OF QUERY]\n",
"What's the weather like in New York in fahrenheit?\n",
"[END OF QUERY]\n",
"\n",
"\n",
"[BEGIN OF HISTORY STEPS]\n",
"[{\"step_id\": 1, \"thought\": \"\", \"tool_calls\": [{\"name\": \"get_weather\", \"arguments\": {\"location\": \"New York\"}}], \"next_observation\": {\"error\": \"\", \"response\": \"Sunny, 81 degrees\"}, \"user_input\": \"Now, search on the Internet for cute puppies\"}]\n",
"[END OF HISTORY STEPS]\n",
"\n"
]
}
],
"source": [
"print(content)"
]
},
{
"cell_type": "markdown",
"id": "71f76a10-a152-49d7-aa6f-3060cc49b935",
"metadata": {},
"source": [
"## Get the model output for follow-up question"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "30af06fd-4aa7-4550-af39-3a77b5951882",
"metadata": {},
"outputs": [],
"source": [
"inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors=\"pt\").to(model.device)\n",
"# 5. Generate the outputs & decode\n",
"# tokenizer.eos_token_id is the id of <|EOT|> token\n",
"outputs = model.generate(inputs, max_new_tokens=512, do_sample=False, num_return_sequences=1, eos_token_id=tokenizer.eos_token_id)\n",
"agent_action = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)\n"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel) (Local)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.13"
}
},
"nbformat": 4,
"nbformat_minor": 5
}