dspy_agent_arxiv_tools_prompt_optimization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimize a tool using Arxiv agent with DSPy and Argilla\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%env OPENAI_API_KEY=<OPEN_AI_API_KEY>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -qqq dspy arxiv argilla ollama"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up DSPy\n",
    "\n",
    "This notebook combines DSPy, Langchain tools, and Argilla to create and optimize an agent for using the Arxiv API. The agent will be able to search for papers, download them, and extract the text. We will start by defining the agent's llm as an openai model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dspy\n",
    "\n",
    "lm = dspy.OllamaLocal(model=\"llama3.2\")\n",
    "# lm = dspy.OpenAI(\n",
    "#     model=\"gpt-4o-mini\",\n",
    "#     api_key=os.getenv(\"OPENAI_API_KEY\"),\n",
    "#     max_tokens=4000,\n",
    "#     temperature=0,\n",
    "# )\n",
    "dspy.settings.configure(lm=lm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining DSPy Signature\n",
    "\n",
    "DSPy relies on signatures to represent data samples. We will define the signature for the Arxiv agent which consists of the following fields:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ArxivQASignature(dspy.Signature):\n",
    "    \"\"\"You will be given a question and an Arxiv Paper ID. Your task is to answer the question.\"\"\"\n",
    "\n",
    "    question: str = dspy.InputField(\n",
    "        prefix=\"Question:\",\n",
    "        desc=\"question to ask\",\n",
    "        format=lambda x: x.strip(),\n",
    "    )\n",
    "    paper_id: str = dspy.InputField(\n",
    "        prefix=\"Paper ID:\",\n",
    "        desc=\"Arxiv Paper ID\",\n",
    "    )\n",
    "    answer: str = dspy.OutputField(\n",
    "        prefix=\"Answer:\",\n",
    "        desc=\"answer to the question\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading the Arxiv Dataset\n",
    "\n",
    "ArXiv QA is a dataset of automated question answering (QA) pairs generated from ArXiv papers using large language models. The dataset includes over 900 papers with corresponding QA pairs, covering a wide range of topics in computer science and related fields. The dataset is organized by year, with papers dating back to 2009."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating train split: 100%|██████████| 210580/210580 [00:00<00:00, 993446.51 examples/s] \n"
     ]
    }
   ],
   "source": [
    "from random import sample\n",
    "from dspy.datasets import DataLoader\n",
    "\n",
    "dl = DataLoader()\n",
    "\n",
    "arxiv_qa = dl.from_huggingface(\n",
    "    \"taesiri/arxiv_qa\",\n",
    "    split=\"train\",\n",
    "    input_keys=(\"question\", \"paper_id\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To keep requests down we'll use a subset of training and testing dataset. We'll be using 100 examples for training set and 20 examples for testing set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "# Set a random seed for reproducibility\n",
    "random.seed(42)\n",
    "\n",
    "aqa_train = [\n",
    "    dspy.Example(\n",
    "        question=example.question, paper_id=example.paper_id, answer=example.answer\n",
    "    ).with_inputs(\"question\", \"paper_id\")\n",
    "    for example in sample(arxiv_qa, 100)\n",
    "]\n",
    "aqa_test = [\n",
    "    dspy.Example(\n",
    "        question=example.question, paper_id=example.paper_id, answer=example.answer\n",
    "    ).with_inputs(\"question\", \"paper_id\")\n",
    "    for example in sample(arxiv_qa, 20)\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DSPy Avatar and Tool\n",
    "\n",
    "We'll setup `Avatar` module with a signature and the tool. Once we have defined our `tools`, we can now create an `Avatar` object by passing the `tools` and `signature`. It takes 2 more optional parameters `verbose` and `max_iters`. `verbose` is used to display the logs and `max_iters` is used to control the number of iterations in multi step execution. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dspy.predict.avatar import Tool, Avatar\n",
    "from langchain_community.utilities import ArxivAPIWrapper\n",
    "\n",
    "tools = [\n",
    "    Tool(\n",
    "        tool=ArxivAPIWrapper(),\n",
    "        name=\"ARXIV_SEARCH\",\n",
    "        desc=\"Pass the arxiv paper id to get the paper information.\",\n",
    "        input_type=\"Arxiv Paper ID\",\n",
    "    ),\n",
    "]\n",
    "\n",
    "arxiv_agent = Avatar(\n",
    "    tools=tools,\n",
    "    signature=ArxivQASignature,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate performance\n",
    "\n",
    "Open enden QA tasks are hard to evaluate on rigid metrics like exact match. So, we'll be using an improvised LLM as Judge for the evaluation of our model on test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Evaluator(dspy.Signature):\n",
    "    \"\"\"Please act as an impartial judge and evaluate the quality of the responses provided by multiple AI assistants to the user question displayed below. You should choose the assistant that offers a better user experience by interacting with the user more effectively and efficiently, and providing a correct final response to the user's question.\n",
    "\n",
    "    Rules:\n",
    "    1. Avoid Position Biases: Ensure that the order in which the responses were presented does not influence your decision. Evaluate each response on its own merits.\n",
    "    2. Length of Responses: Do not let the length of the responses affect your evaluation. Focus on the quality and relevance of the response. A good response is targeted and addresses the user's needs effectively, rather than simply being detailed.\n",
    "    3. Objectivity: Be as objective as possible. Consider the user's perspective and overall experience with each assistant.\"\"\"\n",
    "\n",
    "    question: str = dspy.InputField(\n",
    "        prefix=\"Question:\",\n",
    "        desc=\"question to ask\",\n",
    "    )\n",
    "    reference_answer: str = dspy.InputField(\n",
    "        prefix=\"Reference Answer:\",\n",
    "        desc=\"Answer to the question given by the model.\",\n",
    "    )\n",
    "    answer: str = dspy.InputField(\n",
    "        prefix=\"Answer:\",\n",
    "        desc=\"Answer to the question given by the model.\",\n",
    "    )\n",
    "    rationale: str = dspy.OutputField(\n",
    "        prefix=\"Rationale:\",\n",
    "        desc=\"Explanation of why the answer is correct or incorrect.\",\n",
    "    )\n",
    "    is_correct: bool = dspy.OutputField(\n",
    "        prefix=\"Correct:\",\n",
    "        desc=\"Whether the answer is correct.\",\n",
    "    )\n",
    "\n",
    "\n",
    "evaluator = dspy.TypedPredictor(Evaluator)\n",
    "\n",
    "\n",
    "def metric(example, prediction, trace=None):\n",
    "    return int(\n",
    "        evaluator(\n",
    "            question=example.question,\n",
    "            answer=prediction.answer,\n",
    "            reference_answer=example.answer,\n",
    "        ).is_correct\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For evaluation we can't use `dspy.Evaluate`, reason being that `Avatar` changes it's signature per iteration by adding the actions and it's results to it as fields. So we can create our own hacky thread safe evaluator for it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tqdm\n",
    "\n",
    "from concurrent.futures import ThreadPoolExecutor\n",
    "\n",
    "\n",
    "def process_example(example, signature):\n",
    "    try:\n",
    "        avatar = Avatar(\n",
    "            signature,\n",
    "            tools=tools,\n",
    "            verbose=False,\n",
    "        )\n",
    "        prediction = avatar(**example.inputs().toDict())\n",
    "\n",
    "        return metric(example, prediction)\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        return 0\n",
    "\n",
    "\n",
    "def multi_thread_executor(test_set, signature, num_threads=60):\n",
    "    total_score = 0\n",
    "    total_examples = len(test_set)\n",
    "\n",
    "    with ThreadPoolExecutor(max_workers=num_threads) as executor:\n",
    "        futures = [\n",
    "            executor.submit(process_example, example, signature) for example in test_set\n",
    "        ]\n",
    "\n",
    "        for future in tqdm.tqdm(\n",
    "            futures, total=total_examples, desc=\"Processing examples\"\n",
    "        ):\n",
    "            total_score += future.result()\n",
    "\n",
    "    avg_metric = total_score / total_examples\n",
    "    return avg_metric"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing examples: 100%|██████████| 100/100 [01:53<00:00,  1.14s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Score on ArxivQA: 0.38\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "aqa_score = multi_thread_executor(aqa_test, ArxivQASignature)\n",
    "print(f\"Average Score on ArxivQA: {aqa_score:.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DSPy Optimization with `AvatarOptimizer`\n",
    "\n",
    "To optimize the Actor for optimal tool usage, we use the AvatarOptimizer class. This class is a DSPy implementation of the Avatar method, which uses a comparator module to optimize the Actor for the given tools.\n",
    "\n",
    "Parameters\n",
    "\n",
    "The AvatarOptimizer class takes the following parameters:\n",
    "\n",
    "metric: The metric to optimize for\n",
    "max_iters: The maximum number of iterations for the optimizer\n",
    "lower_bound: The lower bound for the metric to classify an example as negative\n",
    "upper_bound: The upper bound for the metric to classify an example as positive\n",
    "max_positive_inputs: The maximum number of positive inputs to sample for the comparator\n",
    "max_negative_inputs: The maximum number of negative inputs to sample for the comparator\n",
    "optimize_for: Whether to maximize or minimize the metric during optimization\n",
    "Usage\n",
    "\n",
    "To use the AvatarOptimizer, create an instance of the class and pass in the required parameters. Then, call the compile method to optimize the Actor.\n",
    "\n",
    "The AvatarOptimizer optimizes the Actor for optimal tool usage, but it does not optimize the instruction of the signature passed to the Agent. The Actor is the module that directs tool execution and flow, and it is not the same as the signature passed to the Agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dspy.teleprompt import AvatarOptimizer\n",
    "\n",
    "teleprompter = AvatarOptimizer(\n",
    "    metric=metric,\n",
    "    max_iters=1,\n",
    "    max_negative_inputs=10,\n",
    "    max_positive_inputs=10,\n",
    ")\n",
    "\n",
    "optimized_arxiv_agent = teleprompter.compile(student=arxiv_agent, trainset=aqa_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can evaluate our actor module, for this we've provided an implementation of thread safe evaluator that we above as part of class method of `AvatarOptimizer`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "teleprompter.thread_safe_evaluator(aqa_test, optimized_arxiv_agent)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Review optimizations in Argilla\n",
    "\n",
    "Now let's take a look at the opimized tool usage in Argilla. We will pass both agents' responses to the UI and rank them blindly, to see how the optimized agent performs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset(id=UUID('dd04ab88-168d-4bbb-adb9-7f698c58ace2') inserted_at=datetime.datetime(2024, 10, 3, 9, 44, 24, 471390) updated_at=datetime.datetime(2024, 10, 3, 9, 44, 25, 564826) name='arxiv-tools-64c1f828-e922-43d6-91f5-a15902c66ed6' status='ready' guidelines=None allow_extra_metadata=False distribution=OverlapTaskDistributionModel(strategy='overlap', min_submitted=1) workspace_id=UUID('31fb2d4c-46f8-4698-8dcb-27555e78c092') last_activity_at=datetime.datetime(2024, 10, 3, 9, 44, 25, 564826))"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from uuid import uuid4\n",
    "import argilla as rg\n",
    "\n",
    "client = rg.Argilla()\n",
    "\n",
    "dataset = rg.Dataset(\n",
    "    name=f\"arxiv-tools-{uuid4()}\",\n",
    "    settings=rg.Settings(\n",
    "        fields=[\n",
    "            rg.TextField(name=\"response1\"),\n",
    "            rg.TextField(name=\"response2\"),\n",
    "        ],\n",
    "        questions=[\n",
    "            rg.RankingQuestion(name=\"ranking\", values=[\"response1\", \"response2\"])\n",
    "        ],\n",
    "    ),\n",
    ")\n",
    "\n",
    "dataset.create()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's use the agents and push the responses to the UI for ranking."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting the task...\n",
      "Action 1: ARXIV_SEARCH (2403.06404 keywords)\n",
      "Action 2: Finish (The keywords or key terms associated with the paper 2403.06404 are: uncertainty modeling, speaker representation, cosine scoring, neural speaker embedding, embedding estimation, speaker recognition.)\n",
      "Starting the task...\n",
      "Action 1: ARXIV_SEARCH (What are the keywords or key terms associated with the paper 2403.06404?)\n",
      "Action 2: Finish ()\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/ben/code/argilla-cookbook/.venv/lib/python3.12/site-packages/argilla/records/_mapping/_mapper.py:89: UserWarning: Keys ['question'] in data are not present in the mapping and will be ignored.\n",
      "  warnings.warn(f\"Keys {unknown_keys} in data are not present in the mapping and will be ignored.\")\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">DatasetRecords: The provided batch size <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">256</span> was normalized. Using value <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>.\n",
       "</pre>\n"
      ],
      "text/plain": [
       "DatasetRecords: The provided batch size \u001b[1;36m256\u001b[0m was normalized. Using value \u001b[1;36m1\u001b[0m.\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Sending records...: 100%|██████████| 1/1 [00:00<00:00,  1.77batch/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting the task...\n",
      "Action 1: ARXIV_SEARCH (2403.06404 keywords)\n",
      "Action 2: Finish (The keywords or key terms associated with the paper 2403.06404 are: uncertainty modeling, speaker representation, cosine scoring, neural speaker embedding, embedding estimation, speaker recognition.)\n",
      "Starting the task...\n",
      "Action 1: ARXIV_SEARCH (What are the keywords or key terms associated with the paper 2403.06404?)\n",
      "Action 2: Finish ()\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">DatasetRecords: The provided batch size <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">256</span> was normalized. Using value <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>.\n",
       "</pre>\n"
      ],
      "text/plain": [
       "DatasetRecords: The provided batch size \u001b[1;36m256\u001b[0m was normalized. Using value \u001b[1;36m1\u001b[0m.\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Sending records...: 100%|██████████| 1/1 [00:00<00:00,  1.79batch/s]\n"
     ]
    }
   ],
   "source": [
    "max_samples = 2\n",
    "\n",
    "for example in aqa_test[:max_samples]:\n",
    "    question = aqa_test[0].question\n",
    "    base_answer = arxiv_agent(**aqa_test[0]).answer\n",
    "    optimized_answer = optimized_arxiv_agent(**aqa_test[0]).answer\n",
    "    dataset.records.log(\n",
    "        [\n",
    "            {\n",
    "                \"question\": question,\n",
    "                \"response1\": base_answer,\n",
    "                \"response2\": optimized_answer,\n",
    "                \"ranking\": [\"response1\", \"response2\"],\n",
    "            }\n",
    "        ]\n",
    "    )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dspy-74wouE_3-py3.10",
   "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.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}