{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial: Fine-tuning Agents\n",
    "\n",
    "Let's walk through a quick example of optimizing the _language model weights_ (i.e., fine-tuning) inside a DSPy module that represents a ReAct agent playing a game with 50-step tasks.\n",
    "\n",
    "### Install dependencies and download data\n",
    "\n",
    "Install the latest DSPy via `pip install -U dspy` and follow along. This tutorial uses the AlfWorld dataset, which depends on DSPy 2.6.0.\n",
    "\n",
    "You will also need the following dependencies:\n",
    "\n",
    "```shell\n",
    "> pip install -U alfworld==0.3.5 multiprocess\n",
    "> alfworld-download\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<details>\n",
    "<summary>Recommended: Setup MLflow Tracing for learning what's happening under the hood</summary>\n",
    "\n",
    "### MLflow DSPy Integration\n",
    "\n",
    "<a href=\"https://mlflow.org/\">MLflow</a> is an LLMOps tool that natively integrates with DSPy and offer explainability and experiment tracking. In this tutorial, you can use MLflow to visualize prompts and optimization progress as traces to understand the DSPy's behavior better. You can set up MLflow easily by following the four steps below.\n",
    "\n",
    "![MLflow Trace](./mlflow-tracing-agent.png)\n",
    "\n",
    "1. Install MLflow\n",
    "\n",
    "```bash\n",
    "%pip install mlflow>=2.20\n",
    "```\n",
    "\n",
    "2. Start MLflow UI in a separate terminal\n",
    "```bash\n",
    "mlflow ui --port 5000\n",
    "```\n",
    "\n",
    "3. Connect the notebook to MLflow\n",
    "```python\n",
    "import mlflow\n",
    "\n",
    "mlflow.set_tracking_uri(\"http://localhost:5000\")\n",
    "mlflow.set_experiment(\"DSPy\")\n",
    "```\n",
    "\n",
    "4. Enabling tracing.\n",
    "```python\n",
    "mlflow.dspy.autolog()\n",
    "```\n",
    "\n",
    "\n",
    "To learn more about the integration, visit [MLflow DSPy Documentation](https://mlflow.org/docs/latest/llms/dspy/index.html) as well.\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up the language models\n",
    "\n",
    "Our goal is to allow `gpt-4o-mini` to play the AlfWorld household game proficiently, without tinkering with string prompts or example trajectories by hand.\n",
    "\n",
    "Though it's not strictly necessary, we'll make our job a little easier by using the larger `gpt-4o` for prompt optimization and fine-tuning, building our small `gpt-4o-mini` agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dspy\n",
    "\n",
    "gpt4o_mini = dspy.LM('gpt-4o-mini-2024-07-18')\n",
    "gpt4o = dspy.LM('openai/gpt-4o')\n",
    "dspy.configure(experimental=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's load 200 training and 200 development tasks from AlfWorld. The dataset is much larger, but a small number of examples will help keep this tutorial run in 1-2 hours, including fine-tuning.\n",
    "\n",
    "With just 100 training tasks, we'll teach 4o-mini to go from 19% (can barely play the game) to 72%. If you use 500 tasks and retain the demonstrations during fine-tuning, you can push that easily to 82%."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(200, 200)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from dspy.datasets.alfworld import AlfWorld\n",
    "\n",
    "alfworld = AlfWorld()\n",
    "trainset, devset = alfworld.trainset[:200], alfworld.devset[-200:]\n",
    "len(trainset), len(devset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we proceed, let's view an example of this task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-= Welcome to TextWorld, ALFRED! =-\n",
      "\n",
      "You are in the middle of a room. Looking quickly around you, you see a countertop 1, a drawer 8, a drawer 7, a drawer 6, a drawer 5, a drawer 4, a drawer 3, a drawer 2, a drawer 1, a garbagecan 1, a handtowelholder 1, a sinkbasin 2, a sinkbasin 1, a toilet 1, a toiletpaperhanger 1, and a towelholder 1.\n",
      "\n",
      "Your task is to: put a clean soapbar in garbagecan.\n"
     ]
    }
   ],
   "source": [
    "example = trainset[0]\n",
    "\n",
    "with alfworld.POOL.session() as env:\n",
    "    task, info = env.init(**example.inputs())\n",
    "\n",
    "print(task)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Defining the Agent program\n",
    "\n",
    "The agent is a pretty simple `dspy.Module` with one sub-module called `self.react`.\n",
    "\n",
    "This sub-module consumes a definition of a specific `task`, sees its previous `trajectory`, and sees a list of\n",
    "`possible_actions` it can take. It responds simply with the next action.\n",
    "\n",
    "In the `forward` method, we just initialize an environment for the given task `idx`. And we loop up to `self.max_iters`,\n",
    "repeatedly invoking the `self.react` module to take the next action."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Agent(dspy.Module):\n",
    "    def __init__(self, max_iters=50, verbose=False):\n",
    "        self.max_iters = max_iters\n",
    "        self.verbose = verbose\n",
    "        self.react = dspy.Predict(\"task, trajectory, possible_actions: list[str] -> action\")\n",
    "\n",
    "    def forward(self, idx):\n",
    "        with alfworld.POOL.session() as env:\n",
    "            trajectory = []\n",
    "            task, info = env.init(idx)\n",
    "            if self.verbose:\n",
    "                print(f\"Task: {task}\")\n",
    "\n",
    "            for _ in range(self.max_iters):\n",
    "                trajectory_ = \"\\n\".join(trajectory)\n",
    "                possible_actions = info[\"admissible_commands\"][0] + [\"think: ${...thoughts...}\"]\n",
    "                prediction = self.react(task=task, trajectory=trajectory_, possible_actions=possible_actions)\n",
    "                trajectory.append(f\"> {prediction.action}\")\n",
    "\n",
    "                if prediction.action.startswith(\"think:\"):\n",
    "                    trajectory.append(\"OK.\")\n",
    "                    continue\n",
    "\n",
    "                obs, reward, done, info = env.step(prediction.action)\n",
    "                obs, reward, done = obs[0], reward[0], done[0]\n",
    "                trajectory.append(obs)\n",
    "\n",
    "                if self.verbose:\n",
    "                    print(\"\\n\".join(trajectory[-2:]))\n",
    "\n",
    "                if done:\n",
    "                    break\n",
    "\n",
    "        assert reward == int(info[\"won\"][0]), (reward, info[\"won\"][0])\n",
    "        return dspy.Prediction(trajectory=trajectory, success=reward)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Aside: If you wanted to include instructions for your agent...\n",
    "\n",
    "Above, we opted to keep the agent super simple, without even providing short instructions that describe the task.\n",
    "\n",
    "In principle, you can copy a short definition of the AlfWorld task (based on Yao et al., 2022) and use that as the\n",
    "instruction for your agent. This is not inherently essential, but it helps illustrate the role that\n",
    "instructions play in DSPy: they're not for coercing the model to exhibit a certain behavior, but they're there to\n",
    "describe the fundamentals of the task in a straightforward, human-readable way.\n",
    "\n",
    "If you want to do that, you can simply replace this:\n",
    "\n",
    "```python\n",
    "self.react = dspy.Predict(\"task, trajectory, possible_actions: list[str] -> action\")\n",
    "```\n",
    "\n",
    "with this:\n",
    "\n",
    "```python\n",
    "INSTRUCTIONS = \"\"\"\n",
    "Interact with a simulated household to achieve a high-level goal. Make sure to plan, track subgoals,\n",
    "determine likely locations for common household items (e.g. desklamps will likely be on desks, shelfs, or dressers),\n",
    "and explore systematically (e.g. check all desks one by one for desklamp).\n",
    "\"\"\".strip()\n",
    "\n",
    "self.react = dspy.Predict(dspy.Signature(\"task, trajectory, possible_actions: list[str] -> action\", INSTRUCTIONS))\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Zero-shot evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's try this simple program, prior to any optimization work."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Task: -= Welcome to TextWorld, ALFRED! =-\n",
      "\n",
      "You are in the middle of a room. Looking quickly around you, you see a countertop 1, a drawer 8, a drawer 7, a drawer 6, a drawer 5, a drawer 4, a drawer 3, a drawer 2, a drawer 1, a garbagecan 1, a handtowelholder 1, a sinkbasin 2, a sinkbasin 1, a toilet 1, a toiletpaperhanger 1, and a towelholder 1.\n",
      "\n",
      "Your task is to: put a clean soapbar in garbagecan.\n",
      "> go to countertop 1\n",
      "You arrive at countertop 1. On the countertop 1, you see a candle 1, a soapbar 1, a soapbottle 2, a soapbottle 1, and a spraybottle 1.\n",
      "> take soapbar 1 from countertop 1\n",
      "You pick up the soapbar 1 from the countertop 1.\n",
      "> go to garbagecan 1\n",
      "You arrive at garbagecan 1. On the garbagecan 1, you see nothing.\n",
      "> move soapbar 1 to garbagecan 1\n",
      "You move the soapbar 1 to the garbagecan 1.\n",
      "> examine garbagecan 1\n",
      "On the garbagecan 1, you see a soapbar 1.\n",
      "> examine garbagecan 1\n",
      "On the garbagecan 1, you see a soapbar 1.\n",
      "> examine garbagecan 1\n",
      "On the garbagecan 1, you see a soapbar 1.\n",
      "> examine garbagecan 1\n",
      "On the garbagecan 1, you see a soapbar 1.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> take soapbar 1 from garbagecan 1\n",
      "You pick up the soapbar 1 from the garbagecan 1.\n",
      "> move soapbar 1 to garbagecan 1\n",
      "You move the soapbar 1 to the garbagecan 1.\n",
      "> examine garbagecan 1\n",
      "On the garbagecan 1, you see a soapbar 1.\n",
      "> examine garbagecan 1\n",
      "On the garbagecan 1, you see a soapbar 1.\n",
      "> examine garbagecan 1\n",
      "On the garbagecan 1, you see a soapbar 1.\n",
      "> examine garbagecan 1\n",
      "On the garbagecan 1, you see a soapbar 1.\n",
      "> examine garbagecan 1\n",
      "On the garbagecan 1, you see a soapbar 1.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> examine garbagecan 1\n",
      "On the garbagecan 1, you see a soapbar 1.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> examine garbagecan 1\n",
      "On the garbagecan 1, you see a soapbar 1.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n",
      "> examine garbagecan 1\n",
      "On the garbagecan 1, you see a soapbar 1.\n",
      "> look\n",
      "You are facing the garbagecan 1. Next to it, you see nothing.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Prediction(\n",
       "    trajectory=['> go to countertop 1', 'You arrive at countertop 1. On the countertop 1, you see a candle 1, a soapbar 1, a soapbottle 2, a soapbottle 1, and a spraybottle 1.', '> take soapbar 1 from countertop 1', 'You pick up the soapbar 1 from the countertop 1.', '> go to garbagecan 1', 'You arrive at garbagecan 1. On the garbagecan 1, you see nothing.', '> move soapbar 1 to garbagecan 1', 'You move the soapbar 1 to the garbagecan 1.', '> examine garbagecan 1', 'On the garbagecan 1, you see a soapbar 1.', '> examine garbagecan 1', 'On the garbagecan 1, you see a soapbar 1.', '> examine garbagecan 1', 'On the garbagecan 1, you see a soapbar 1.', '> examine garbagecan 1', 'On the garbagecan 1, you see a soapbar 1.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> take soapbar 1 from garbagecan 1', 'You pick up the soapbar 1 from the garbagecan 1.', '> move soapbar 1 to garbagecan 1', 'You move the soapbar 1 to the garbagecan 1.', '> examine garbagecan 1', 'On the garbagecan 1, you see a soapbar 1.', '> examine garbagecan 1', 'On the garbagecan 1, you see a soapbar 1.', '> examine garbagecan 1', 'On the garbagecan 1, you see a soapbar 1.', '> examine garbagecan 1', 'On the garbagecan 1, you see a soapbar 1.', '> examine garbagecan 1', 'On the garbagecan 1, you see a soapbar 1.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> examine garbagecan 1', 'On the garbagecan 1, you see a soapbar 1.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> examine garbagecan 1', 'On the garbagecan 1, you see a soapbar 1.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.', '> examine garbagecan 1', 'On the garbagecan 1, you see a soapbar 1.', '> look', 'You are facing the garbagecan 1. Next to it, you see nothing.'],\n",
       "    success=0\n",
       ")"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_4o = Agent()\n",
    "agent_4o.set_lm(gpt4o)\n",
    "agent_4o.verbose = True\n",
    "\n",
    "agent_4o(**example.inputs())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Okay, in this case it couldn't solve this example! Now, let's check the average quality of 4o and 4o-mini."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "metric = lambda x, y, trace=None: y.success\n",
    "evaluate = dspy.Evaluate(devset=devset, metric=metric, display_progress=True, num_threads=16)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<details>\n",
    "<summary>Tracking Evaluation Results in MLflow Experiment</summary>\n",
    "\n",
    "<br/>\n",
    "\n",
    "To track and visualize the evaluation results over time, you can record the results in MLflow Experiment.\n",
    "\n",
    "\n",
    "```python\n",
    "import mlflow\n",
    "\n",
    "with mlflow.start_run(run_name=\"agent_evaluation\"):\n",
    "    evaluate = dspy.Evaluate(\n",
    "        devset=devset,\n",
    "        metric=metric,\n",
    "        num_threads=16,\n",
    "        display_progress=True,\n",
    "        # To record the outputs and detailed scores to MLflow\n",
    "        return_all_scores=True,\n",
    "        return_outputs=True,\n",
    "    )\n",
    "\n",
    "    # Evaluate the program as usual\n",
    "    aggregated_score, outputs, all_scores = evaluate(cot)\n",
    "\n",
    "    # Log the aggregated score\n",
    "    mlflow.log_metric(\"success_rate\", aggregated_score)\n",
    "    # Log the detailed evaluation results as a table\n",
    "    mlflow.log_table(\n",
    "        {\n",
    "            \"Idx\": [example.idx for example in eval_set],\n",
    "            \"Result\": outputs,\n",
    "            \"Success\": all_scores,\n",
    "        },\n",
    "        artifact_file=\"eval_results.json\",\n",
    "    )\n",
    "```\n",
    "\n",
    "To learn more about the integration, visit [MLflow DSPy Documentation](https://mlflow.org/docs/latest/llms/dspy/index.html) as well.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Metric: 115.00 / 200 (57.5%): 100%|██████████| 200/200 [06:14<00:00,  1.87s/it]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024/12/28 11:10:25 INFO dspy.evaluate.evaluate: Average Metric: 115 / 200 (57.5%)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "57.5"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_4o.verbose = False\n",
    "evaluate(agent_4o)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Metric: 30.00 / 200 (15.0%): 100%|██████████| 200/200 [08:33<00:00,  2.57s/it]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024/12/28 11:18:59 INFO dspy.evaluate.evaluate: Average Metric: 30 / 200 (15.0%)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "15.0"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_4o_mini = Agent()\n",
    "agent_4o_mini.set_lm(gpt4o_mini)\n",
    "\n",
    "evaluate(agent_4o_mini)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Out of the box, on this task, 4o is decent (58% success rate) while 4o-mini struggles (15% success rate).\n",
    "\n",
    "Let's apply the following strategy:\n",
    "\n",
    "1. We'll optimize the _prompts_ for gpt-4o in a lightweight way.\n",
    "2. We'll then use this prompt-optimized agent as a teacher to fine-tune gpt-4o-mini on the task. This will increase its quality from 19% to 72% (or 82% if you use 500 trainset examples)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prompt-optimizing GPT-4o"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = dspy.MIPROv2(metric=metric, auto=\"light\", num_threads=16, prompt_model=gpt4o)\n",
    "\n",
    "config = dict(max_bootstrapped_demos=1, max_labeled_demos=0, minibatch_size=40)\n",
    "optimized_4o = optimizer.compile(agent_4o, trainset=trainset, **config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fine-tuning GPT-4o-mini\n",
    "\n",
    "For fine-tuning, we'll need a teacher program (`optimized_4o` above) and a student program derived from it (`student_4om` below)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "student_4o_mini = optimized_4o.deepcopy()\n",
    "student_4o_mini.set_lm(gpt4o_mini)\n",
    "# student_4o_mini.react.demos = []  # you can optionally reset the demos"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = dspy.BootstrapFinetune(metric=metric, num_threads=16)\n",
    "finetuned_4o_mini = optimizer.compile(student_4o_mini, teacher=optimized_4o, trainset=trainset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluate the finetuned GPT-4o-mini agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Metric: 143.00 / 200 (71.5%): 100%|██████████| 200/200 [03:15<00:00,  1.05it/s]"
     ]
    }
   ],
   "source": [
    "evaluate(finetuned_4o_mini)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Having done all this optimization, let's save our program so we can use it later! This will keep a reference to the fine-tuned model as well, as long as it continued to exist with the same identifier at the provider side."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "finetuned_4o_mini.save('finetuned_4o_mini_001.pkl')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<details>\n",
    "<summary>Saving programs in MLflow Experiment</summary>\n",
    "\n",
    "<br/>\n",
    "\n",
    "Instead of saving the program to a local file, you can track it in MLflow for better reproducibility and collaboration.\n",
    "\n",
    "1. **Dependency Management**: MLflow automatically save the frozen environment metadata along with the program to ensure reproducibility.\n",
    "2. **Experiment Tracking**: With MLflow, you can track the program's performance and cost along with the program itself.\n",
    "3. **Collaboration**: You can share the program and results with your team members by sharing the MLflow experiment.\n",
    "\n",
    "To save the program in MLflow, run the following code:\n",
    "\n",
    "```python\n",
    "import mlflow\n",
    "\n",
    "# Start an MLflow Run and save the program\n",
    "with mlflow.start_run(run_name=\"optimized\"):\n",
    "    model_info = mlflow.dspy.log_model(\n",
    "        finetuned_4o_mini,\n",
    "        artifact_path=\"model\", # Any name to save the program in MLflow\n",
    "    )\n",
    "\n",
    "# Load the program back from MLflow\n",
    "loaded = mlflow.dspy.load_model(model_info.model_uri)\n",
    "```\n",
    "\n",
    "To learn more about the integration, visit [MLflow DSPy Documentation](https://mlflow.org/docs/latest/llms/dspy/index.html) as well.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now check out one task using our fine-tuned agent program!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Task: -= Welcome to TextWorld, ALFRED! =-\n",
      "\n",
      "You are in the middle of a room. Looking quickly around you, you see a armchair 1, a cabinet 1, a drawer 21, a drawer 20, a drawer 19, a drawer 18, a drawer 17, a drawer 16, a drawer 15, a drawer 14, a drawer 13, a drawer 12, a drawer 11, a drawer 10, a drawer 9, a drawer 8, a drawer 7, a drawer 6, a drawer 5, a drawer 4, a drawer 3, a drawer 2, a drawer 1, a dresser 1, a garbagecan 1, a sidetable 5, a sidetable 4, a sidetable 3, a sidetable 2, a sidetable 1, and a sofa 1.\n",
      "\n",
      "Your task is to: put some box on dresser.\n",
      "> go to cabinet 1\n",
      "You arrive at cabinet 1. On the cabinet 1, you see nothing.\n",
      "> go to dresser 1\n",
      "You arrive at dresser 1. On the dresser 1, you see a book 1, a newspaper 1, a remotecontrol 1, a statue 3, and a television 1.\n",
      "> look\n",
      "You are facing the dresser 1. Next to it, you see nothing.\n",
      "> go to sidetable 1\n",
      "You arrive at sidetable 1. On the sidetable 1, you see a cellphone 1, and a desklamp 1.\n",
      "> go to sidetable 2\n",
      "You arrive at sidetable 2. On the sidetable 2, you see a box 2.\n",
      "> take box 2 from sidetable 2\n",
      "You pick up the box 2 from the sidetable 2.\n",
      "> go to dresser 1\n",
      "You arrive at dresser 1. On the dresser 1, you see a book 1, a newspaper 1, a remotecontrol 1, a statue 3, and a television 1.\n",
      "> move box 2 to dresser 1\n",
      "You move the box 2 to the dresser 1.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Prediction(\n",
       "    trajectory=['> go to cabinet 1', 'You arrive at cabinet 1. On the cabinet 1, you see nothing.', '> go to dresser 1', 'You arrive at dresser 1. On the dresser 1, you see a book 1, a newspaper 1, a remotecontrol 1, a statue 3, and a television 1.', '> look', 'You are facing the dresser 1. Next to it, you see nothing.', '> go to sidetable 1', 'You arrive at sidetable 1. On the sidetable 1, you see a cellphone 1, and a desklamp 1.', '> go to sidetable 2', 'You arrive at sidetable 2. On the sidetable 2, you see a box 2.', '> take box 2 from sidetable 2', 'You pick up the box 2 from the sidetable 2.', '> go to dresser 1', 'You arrive at dresser 1. On the dresser 1, you see a book 1, a newspaper 1, a remotecontrol 1, a statue 3, and a television 1.', '> move box 2 to dresser 1', 'You move the box 2 to the dresser 1.'],\n",
       "    success=1\n",
       ")"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "finetuned_4o_mini.verbose = True\n",
    "finetuned_4o_mini(**devset[0].inputs())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you want to load and use the agent program, you can do that as follows."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **⚠️ Security Warning:** Loading `.pkl` files can execute arbitrary code and may be dangerous. Only save and load pickle files from trusted sources in secure environments. Consider using JSON format when possible for safer serialization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "loaded = Agent()\n",
    "loaded.load('finetuned_4o_mini_001.pkl', allow_pickle=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "jun2024_py310",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
