From 2c208e3c50bebfd5f746b78d3daaaf415aee249c Mon Sep 17 00:00:00 2001 From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com> Date: Mon, 19 Feb 2024 21:08:45 -0800 Subject: [PATCH] tmp --- examples/lats/lats.ipynb | 491 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 491 insertions(+) create mode 100644 examples/lats/lats.ipynb diff --git a/examples/lats/lats.ipynb b/examples/lats/lats.ipynb new file mode 100644 index 000000000..0f3719235 --- /dev/null +++ b/examples/lats/lats.ipynb @@ -0,0 +1,491 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "9e0c1743-8775-4de2-a599-78b050551489", + "metadata": {}, + "source": [ + "# Language Agent Tree Search\n", + "\n", + "\n", + "The basic idea is to combine reflection/evaluation and search (specifically monte-carlo trees search) to get better overall performance.\n", + "\n", + "As with all search-based methods, you trade off some inference-time compute and execution time for quality, and you can use langsmith them to extract the best trajectory for finetuning to reduce the need for as expensive search in the future..\n", + "\n", + "#### Expand and simulate\n", + "1. Execute N actions/tasks\n", + "2. Get observations from tasks\n", + "3. Evaluate the response -> give value of action V\n", + "4. Update global state with the next step from this rollout\n", + "#### Reflect\n", + "For all terminal states (final responses), \n", + "1. Get the reward\n", + "2. Assign `c` = reflection IFF reward isn't a 'sucess'\n", + "#### Select\n", + "(from (Upper Confidence bounds applied to Trees))\n", + "1. next action is the one that maximizes the value of V + w sqrt(ln (counts(prev state) )/counts(current state))\n", + "\n", + "\n", + "#### \"Backgpropagate\"\n", + "\n", + "```\n", + "V(s) = (V_{old}(s)(N(s) - 1)+r) / N(s)\n", + "```\n", + "Nodes in the path are equivalent to the trajectories.\n", + "\n", + "Then after the backpropagation, we pick the " + ] + }, + { + "cell_type": "markdown", + "id": "db28668b-5491-4c93-a961-bd339f09202c", + "metadata": {}, + "source": [ + "## 0. Prerequisites\n", + "\n", + "Install `langgraph` (for the framework), `langchain_openai` (for the LLM), and `langchain` + `tavily-python` (for the search engine).\n", + "\n", + "We will use tavily search as a tool. You can get an API key [here](https://app.tavily.com/sign-in) or replace with a different tool of your choosing." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "dcc9159b-cc8c-426d-9670-3e8ada06723f", + "metadata": {}, + "outputs": [], + "source": [ + "# %pip install -U --quiet langchain langgraph langchain_openai\n", + "# %pip install -U --quiet tavily-python" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "5f43795e-353d-4a35-8dc9-c867b0d18568", + "metadata": {}, + "outputs": [], + "source": [ + "import getpass\n", + "import os\n", + "\n", + "\n", + "def _set_if_undefined(var: str) -> None:\n", + " if os.environ.get(var):\n", + " return\n", + " os.environ[var] = getpass.getpass(var)\n", + "\n", + "\n", + "# Optional: Configure tracing to visualize and debug the agent\n", + "_set_if_undefined(\"LANGCHAIN_API_KEY\")\n", + "os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n", + "os.environ[\"LANGCHAIN_PROJECT\"] = \"LATS\"\n", + "\n", + "_set_if_undefined(\"OPENAI_API_KEY\")\n", + "_set_if_undefined(\"TAVILY_API_KEY\")" + ] + }, + { + "cell_type": "markdown", + "id": "f857eacb-af4a-47d1-b45f-da74941125c2", + "metadata": {}, + "source": [ + "## Minimal example\n", + "\n", + "Either do code generation or search. Maybe let's just to web research since we don't want to do full code rollouts." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "07f0c534-598b-4ce6-8ec3-eba1e9241a88", + "metadata": {}, + "outputs": [], + "source": [ + "from langchain_community.tools.tavily_search import TavilySearchResults\n", + "from langchain_community.utilities.tavily_search import TavilySearchAPIWrapper\n", + "\n", + "search = TavilySearchAPIWrapper()\n", + "tavily_tool = TavilySearchResults(api_wrapper=search, max_results=5)" + ] + }, + { + "cell_type": "code", + "execution_count": 21, + "id": "1ff6248c-a7b7-4be1-9c5e-576a05d719ed", + "metadata": {}, + "outputs": [], + "source": [ + "from collections import defaultdict\n", + "from typing import List\n", + "\n", + "from langchain.output_parsers.openai_tools import (\n", + " JsonOutputToolsParser,\n", + " PydanticToolsParser,\n", + ")\n", + "from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, ToolMessage\n", + "\n", + "from langgraph.prebuilt.tool_executor import ToolExecutor, ToolInvocation" + ] + }, + { + "cell_type": "code", + "execution_count": 160, + "id": "72fc5363-f0f3-4362-8499-14eb583bd75b", + "metadata": {}, + "outputs": [], + "source": [ + "import datetime\n", + "from typing import List\n", + "\n", + "from langchain.output_parsers.openai_tools import JsonOutputToolsParser\n", + "from langchain_core.output_parsers import StrOutputParser\n", + "from langchain_core.prompt_values import PromptValue\n", + "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n", + "from langchain_core.pydantic_v1 import BaseModel, Field, ValidationError\n", + "from langchain_openai import ChatOpenAI\n", + "\n", + "prompt_template = ChatPromptTemplate.from_messages(\n", + " [\n", + " (\n", + " \"system\",\n", + " \"You are a helpful assistant.\",\n", + " ),\n", + " (\"user\", \"{input}\"),\n", + " MessagesPlaceholder(variable_name=\"messages\", optional=True),\n", + " ]\n", + ")\n", + "\n", + "\n", + "llm = ChatOpenAI(model=\"gpt-3.5-turbo\")\n", + "initial_answer_chain = prompt_template | llm | StrOutputParser()\n", + "\n", + "\n", + "def expand(prompt_value: PromptValue, config):\n", + " n = config[\"configurable\"].get(\"n\", 5)\n", + " llm_res = llm.generate([prompt_value.to_messages()], n=n, temperature=1.0)\n", + " # Could consider scoring dupped values higher\n", + " dedupped = {gen.text: gen for gen in llm_res.generations[0]}\n", + " return [AIMessage(content=gen.text) for gen in dedupped.values()]\n", + "\n", + "\n", + "expansion_prompt_template = ChatPromptTemplate.from_messages(\n", + " [\n", + " (\n", + " \"system\",\n", + " \"You are a helpful assistant. Generate an \"\n", + " \"improved response based on the provided reflections.\",\n", + " ),\n", + " (\"user\", \"{input}\"),\n", + " MessagesPlaceholder(variable_name=\"messages\", optional=True),\n", + " ]\n", + ")\n", + "expansion_chain = expansion_prompt_template | expand" + ] + }, + { + "cell_type": "code", + "execution_count": 148, + "id": "8a215bc6-f3fc-4b4d-98b1-eb760e4c9199", + "metadata": {}, + "outputs": [], + "source": [ + "from langchain_core.pydantic_v1 import BaseModel\n", + "\n", + "\n", + "class Reflection(BaseModel):\n", + " critique: str = Field(\n", + " description=\"~50-100 word critique of the current response, \"\n", + " \"outline parts that are missing, superfluous, or low quality.\"\n", + " )\n", + " is_finished: bool = Field(\n", + " description=\"Whether this response completely and accurately resolves the user's request.\"\n", + " )\n", + " score: int = Field(gte=1, lte=10, description=\"Score of the response.\")" + ] + }, + { + "cell_type": "code", + "execution_count": 164, + "id": "c1219cd1-6cfb-4c83-8579-d3a003569150", + "metadata": {}, + "outputs": [], + "source": [ + "from typing import Optional\n", + "\n", + "\n", + "class Node:\n", + " def __init__(\n", + " self,\n", + " solution: str,\n", + " parent: Optional[Node] = None,\n", + " reflection: Optional[Reflection] = None,\n", + " ):\n", + " self.solution = solution\n", + " self.parent = parent\n", + " self.children = []\n", + " self.value = 0\n", + " self.visits = 0\n", + " self.reflection = reflection\n", + " self.is_finished = False\n", + " if reflection is None and reflection.is_finished:\n", + " self._mark_as_finished()\n", + "\n", + " def upper_confidence_bound(self, exploration_weight: float = 1.0):\n", + " if self.parent is None:\n", + " raise ValueError(\"Cannot compute UTC for root node.\")\n", + " if self.visits == 0:\n", + " return self.value\n", + " return (self.value / self.visits) + exploration_weight * math.sqrt(\n", + " math.log(self.parent.visits) / self.visits\n", + " )\n", + "\n", + " @property\n", + " def best_child(self):\n", + " if not self.children:\n", + " return None\n", + " return max(self.children, key=lambda child: child.uct())\n", + "\n", + " @property\n", + " def best_child_score(self):\n", + " if not self.children:\n", + " return None\n", + " return max(self.children, key=lambda child: child.value)\n", + "\n", + " @property\n", + " def depth_at_node(self):\n", + " if self.children:\n", + " return 1 + max([child.depth_at_node for child in self.children])\n", + " return 1\n", + "\n", + " def add_reflection(self, reflection: Reflection):\n", + " if self.reflection is not None:\n", + " raise ValueError(\"Cannot overwrite existing reflection\")\n", + " self.reflection = reflection\n", + " if reflection is None and reflection.is_finished:\n", + " self._mark_as_finished()\n", + "\n", + " def update(self, reward: float):\n", + " self.visits += 1\n", + " self.value += reward\n", + "\n", + " def _mark_as_finished(self):\n", + " self.is_finished = True\n", + " parent = self.parent\n", + " while parent:\n", + " parent.is_finished = True\n", + " parent = parent.parent" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c01ae663-56fe-40fc-bac2-29f0bc7765a5", + "metadata": {}, + "outputs": [], + "source": [ + "n=5 # number of generated actions (beam width)\n", + "w=1 # Exploration weight. \n", + "L = 5 # Depth limit / max rollout\n", + "K = 3 # Number of rollouts\n", + "c.. # context" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "63ef2b7d-d73e-4073-b549-f37b96011145", + "metadata": {}, + "outputs": [], + "source": [ + "from " + ] + }, + { + "cell_type": "markdown", + "id": "d9c1ba9b-a083-443c-af82-fa116fc450de", + "metadata": {}, + "source": [ + "```\n", + "State:\n", + "- root\n", + "\n", + "Nodes:\n", + "# first_attempt: llm -> generate code\n", + "# Rollout\n", + " - Selection\n", + " - Expansion\n", + "# Reflection / Simulation\n", + "# Backpropagation\n", + "```" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bc7ba5b1-0f4e-40cb-a4d1-f86dc935742f", + "metadata": {}, + "outputs": [], + "source": [ + "# with langsmith.trace(\"mcts\", inputs={\"item\": item}) as trace:\n", + "# # Get the first output\n", + "# cur_func_impl = gen.func_impl(item[\"prompt\"], model, \"simple\")\n", + "# root = Node(cur_func_impl) # initial solution (for pass@1 metric)\n", + "# test_feedback.append(feedback)\n", + "# with langsmith.trace(\n", + "# \"self_reflection\", inputs={\"fun\": cur_func_impl}\n", + "# ) as trace:\n", + "# reflection = gen.self_reflection(cur_func_impl, feedback, model)\n", + "# trace.end(outputs={\"reflection\": reflection})\n", + "# reflections += [reflection]\n", + "# root.test_feedback = feedback\n", + "# root.reflection = reflection\n", + "\n", + "# for cur_iter in range(max_iters):\n", + "# # Selection\n", + "\n", + "# node = root\n", + "# trajectory = {\"solutions\": [], \"feedbacks\": []}\n", + "\n", + "# while node.children:\n", + "# node = node.best_child()\n", + "# trajectory[\"solutions\"].append(node.solution)\n", + "\n", + "# # Expansion\n", + "# for _ in range(n):\n", + "# new_solution = None\n", + "# strategy = \"mcts\"\n", + "# prev_func_impl = node.solution\n", + "# feedback = node.test_feedback\n", + "# reflection = node.reflection\n", + "# acc_feedback, acc_reflection = gather_context_from_tree(node)\n", + "# new_solution = gen.func_impl(\n", + "# func_sig=item[\"prompt\"],\n", + "# model=model,\n", + "# strategy=strategy,\n", + "# prev_func_impl=prev_func_impl,\n", + "# feedback=feedback,\n", + "# self_reflection=reflection,\n", + "# acc_feedback=acc_feedback,\n", + "# acc_reflection=acc_reflection,\n", + "# )\n", + "\n", + "# combined_context = \"\\nPrevious Trial\\n\\n\" + new_solution\n", + "\n", + "# child = Node(\n", + "# new_solution,\n", + "# parent=node,\n", + "# context=combined_context,\n", + "# depth=node.depth + 1,\n", + "# )\n", + "# node.children.append(child)\n", + "\n", + "# # Simulation\n", + "# reward_real = 0\n", + "# for child in node.children:\n", + "# reflection = gen.self_reflection(\n", + "# child.solution, feedback_internal, model\n", + "# )\n", + "# reflections.append(reflection)\n", + "# child.reflection = reflection\n", + "# child.test_feedback = feedback_internal\n", + "# child.context += (\n", + "# \"\\n\\nPrevious Trial\\n\\n\"\n", + "# + child.solution\n", + "# + \"\\n\\nTest results: \\n\"\n", + "# + feedback_internal\n", + "# + \"\\n\\nSelf-reflection: \"\n", + "# + reflection\n", + "# )\n", + "\n", + "# if \"Tested passed:\" in feedback_internal:\n", + "# # Split at \"Tests failed:\" and get the part before it (which contains the passed tests)\n", + "# passed_section = feedback_internal.split(\"Tests failed:\")[0]\n", + "# # Split at \"Tested passed:\" and get the part after it, then count the non-empty lines\n", + "# reward_internal = len(\n", + "# [\n", + "# line\n", + "# for line in passed_section.split(\"Tested passed:\")[\n", + "# 1\n", + "# ].splitlines()\n", + "# if line.strip() != \"\"\n", + "# ]\n", + "# )\n", + "# reward_internal = reward_internal / len(tests_i)\n", + "# else:\n", + "# reward_internal = 0\n", + "# if is_passing_internal or cur_iter == max_iters - 1:\n", + "# is_passing = exe.evaluate(\n", + "# item[\"entry_point\"],\n", + "# child.solution,\n", + "# item[\"test\"],\n", + "# timeout=10,\n", + "# )\n", + "# if is_passing:\n", + "# item[\"solution\"] = child.solution\n", + "# is_solved = True\n", + "# reward_real = 1\n", + "# break\n", + "\n", + "# if is_solved:\n", + "# break\n", + "# reward = reward_internal + reward_real\n", + "# child.update(reward)\n", + "\n", + "# # Backpropagation\n", + "# temp = child\n", + "# while temp.parent:\n", + "# temp = temp.parent\n", + "# temp.update(reward)\n", + "\n", + "# # Choose the best solution after all iterations\n", + "# best_solution = root.best_child_value().solution\n", + "# item[\"solution\"] = best_solution\n", + "# reflections.append(\"MCTS reflections\")\n", + "# implementations.append(best_solution)" + ] + }, + { + "cell_type": "code", + "execution_count": 144, + "id": "8ebb3a4f-d81e-4238-9c60-8004cedef454", + "metadata": {}, + "outputs": [], + "source": [ + "# from datasets import load_dataset\n", + "\n", + "# dataset = load_dataset(\"deepmind/code_contests\", split=\"valid\")\n", + "# example = dataset[70]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9f4525b9-f63e-449d-9816-e37ca726a522", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "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.11.2" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}