Upload multi_turn_xlam.ipynb

example/multi_turn_xlam.ipynb

@@ -0,0 +1,459 @@
+{
+ "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
+}