{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Tutorial: Entity Extraction\n",
        "\n",
        "This tutorial demonstrates how to perform **entity extraction** using the CoNLL-2003 dataset with DSPy. The focus is on extracting entities referring to people. We will:\n",
        "\n",
        "- Extract and label entities from the CoNLL-2003 dataset that refer to people\n",
        "- Define a DSPy program for extracting entities that refer to people\n",
        "- Optimize and evaluate the program on a subset of the CoNLL-2003 dataset\n",
        "\n",
        "By the end of this tutorial, you'll understand how to structure tasks in DSPy using signatures and modules, evaluate your system's performance, and improve its quality with optimizers.\n",
        "\n",
        "Install the latest version of DSPy and follow along. If you're looking instead for a conceptual overview of DSPy, this [recent lecture](https://www.youtube.com/live/JEMYuzrKLUw) is a good place to start."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Install the latest version of DSPy\n",
        "%pip install -U dspy\n",
        "# Install the Hugging Face datasets library to load the CoNLL-2003 dataset\n",
        "%pip install datasets"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<details>\n",
        "<summary>Recommended: Set up MLflow Tracing to understand 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-entity-extraction.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",
        "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": [
        "## Load and Prepare the Dataset\n",
        "\n",
        "In this section, we prepare the CoNLL-2003 dataset, which is commonly used for entity extraction tasks. The dataset includes tokens annotated with entity labels such as persons, organizations, and locations.\n",
        "\n",
        "We will:\n",
        "1. Load the dataset using the Hugging Face `datasets` library.\n",
        "2. Define a function to extract tokens referring to people.\n",
        "3. Slice the dataset to create smaller subsets for training and testing.\n",
        "\n",
        "DSPy expects examples in a structured format, so we'll also transform the dataset into DSPy `Examples` for easy integration."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [],
      "source": [
        "import os\n",
        "import tempfile\n",
        "from datasets import load_dataset\n",
        "from typing import Dict, Any, List\n",
        "import dspy\n",
        "\n",
        "def load_conll_dataset() -> dict:\n",
        "    \"\"\"\n",
        "    Loads the CoNLL-2003 dataset into train, validation, and test splits.\n",
        "    \n",
        "    Returns:\n",
        "        dict: Dataset splits with keys 'train', 'validation', and 'test'.\n",
        "    \"\"\"\n",
        "    with tempfile.TemporaryDirectory() as temp_dir:\n",
        "        # Use a temporary Hugging Face cache directory for compatibility with certain hosted notebook\n",
        "        # environments that don't support the default Hugging Face cache directory\n",
        "        os.environ[\"HF_DATASETS_CACHE\"] = temp_dir\n",
        "        return load_dataset(\"lhoestq/conll2003\")\n",
        "\n",
        "def extract_people_entities(data_row: dict[str, Any]) -> list[str]:\n",
        "    \"\"\"\n",
        "    Extracts entities referring to people from a row of the CoNLL-2003 dataset.\n",
        "    \n",
        "    Args:\n",
        "        data_row (dict[str, Any]): A row from the dataset containing tokens and NER tags.\n",
        "    \n",
        "    Returns:\n",
        "        list[str]: List of tokens tagged as people.\n",
        "    \"\"\"\n",
        "    return [\n",
        "        token\n",
        "        for token, ner_tag in zip(data_row[\"tokens\"], data_row[\"ner_tags\"])\n",
        "        if ner_tag in (1, 2)  # CoNLL entity codes 1 and 2 refer to people\n",
        "    ]\n",
        "\n",
        "def prepare_dataset(data_split, start: int, end: int) -> list[dspy.Example]:\n",
        "    \"\"\"\n",
        "    Prepares a sliced dataset split for use with DSPy.\n",
        "    \n",
        "    Args:\n",
        "        data_split: The dataset split (e.g., train or test).\n",
        "        start (int): Starting index of the slice.\n",
        "        end (int): Ending index of the slice.\n",
        "    \n",
        "    Returns:\n",
        "        list[dspy.Example]: List of DSPy Examples with tokens and expected labels.\n",
        "    \"\"\"\n",
        "    return [\n",
        "        dspy.Example(\n",
        "            tokens=row[\"tokens\"],\n",
        "            expected_extracted_people=extract_people_entities(row)\n",
        "        ).with_inputs(\"tokens\")\n",
        "        for row in data_split.select(range(start, end))\n",
        "    ]\n",
        "\n",
        "# Load the dataset\n",
        "dataset = load_conll_dataset()\n",
        "\n",
        "# Prepare the training and test sets\n",
        "train_set = prepare_dataset(dataset[\"train\"], 0, 50)\n",
        "test_set = prepare_dataset(dataset[\"test\"], 0, 200)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Configure DSPy and create an Entity Extraction Program\n",
        "\n",
        "Here, we define a DSPy program for extracting entities referring to people from tokenized text.\n",
        "\n",
        "Then, we configure DSPy to use a particular language model (`gpt-4o-mini`) for all invocations of the program.\n",
        "\n",
        "**Key DSPy Concepts Introduced:**\n",
        "- **Signatures:** Define structured input/output schemas for your program.\n",
        "- **Modules:** Encapsulate program logic in reusable, composable units.\n",
        "\n",
        "Specifically, we'll:\n",
        "- Create a `PeopleExtraction` DSPy Signature to specify the input (`tokens`) and output (`extracted_people`) fields.\n",
        "- Define a `people_extractor` program that uses DSPy's built-in `dspy.ChainOfThought` module to implement the `PeopleExtraction` signature. The program extracts entities referring to people from a list of input tokens using language model (LM) prompting.\n",
        "- Use the `dspy.LM` class and `dspy.configure()` method to configure the language model that DSPy will use when invoking the program."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [],
      "source": [
        "from typing import List\n",
        "\n",
        "class PeopleExtraction(dspy.Signature):\n",
        "    \"\"\"\n",
        "    Extract contiguous tokens referring to specific people, if any, from a list of string tokens.\n",
        "    Output a list of tokens. In other words, do not combine multiple tokens into a single value.\n",
        "    \"\"\"\n",
        "    tokens: list[str] = dspy.InputField(desc=\"tokenized text\")\n",
        "    extracted_people: list[str] = dspy.OutputField(desc=\"all tokens referring to specific people extracted from the tokenized text\")\n",
        "\n",
        "people_extractor = dspy.ChainOfThought(PeopleExtraction)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Here, we tell DSPy to use OpenAI's `gpt-4o-mini` model in our program. To authenticate, DSPy reads your `OPENAI_API_KEY`. You can easily swap this out for [other providers or local models](https://github.com/stanfordnlp/dspy/blob/main/examples/migration.ipynb)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [],
      "source": [
        "lm = dspy.LM(model=\"openai/gpt-4o-mini\")\n",
        "dspy.configure(lm=lm)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Define Metric and Evaluation Functions\n",
        "\n",
        "In DSPy, evaluating a program's performance is critical for iterative development. A good evaluation framework allows us to:\n",
        "- Measure the quality of our program's outputs.\n",
        "- Compare outputs against ground-truth labels.\n",
        "- Identify areas for improvement.\n",
        "\n",
        "**What We'll Do:**\n",
        "- Define a custom metric (`extraction_correctness_metric`) to evaluate whether the extracted entities match the ground truth.\n",
        "- Create an evaluation function (`evaluate_correctness`) to apply this metric to a training or test dataset and compute the overall accuracy.\n",
        "\n",
        "The evaluation function uses DSPy's `Evaluate` utility to handle parallelism and visualization of results."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [],
      "source": [
        "def extraction_correctness_metric(example: dspy.Example, prediction: dspy.Prediction, trace=None) -> bool:\n",
        "    \"\"\"\n",
        "    Computes correctness of entity extraction predictions.\n",
        "    \n",
        "    Args:\n",
        "        example (dspy.Example): The dataset example containing expected people entities.\n",
        "        prediction (dspy.Prediction): The prediction from the DSPy people extraction program.\n",
        "        trace: Optional trace object for debugging.\n",
        "    \n",
        "    Returns:\n",
        "        bool: True if predictions match expectations, False otherwise.\n",
        "    \"\"\"\n",
        "    return prediction.extracted_people == example.expected_extracted_people\n",
        "\n",
        "evaluate_correctness = dspy.Evaluate(\n",
        "    devset=test_set,\n",
        "    metric=extraction_correctness_metric,\n",
        "    num_threads=24,\n",
        "    display_progress=True,\n",
        "    display_table=True\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Evaluate Initial Extractor\n",
        "\n",
        "Before optimizing our program, we need a baseline evaluation to understand its current performance. This helps us:\n",
        "- Establish a reference point for comparison after optimization.\n",
        "- Identify potential weaknesses in the initial implementation.\n",
        "\n",
        "In this step, we'll run our `people_extractor` program on the test set and measure its accuracy using the evaluation framework defined earlier."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Average Metric: 172.00 / 200 (86.0%): 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:16<00:00, 11.94it/s]"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2024/11/18 21:08:04 INFO dspy.evaluate.evaluate: Average Metric: 172 / 200 (86.0%)\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>tokens</th>\n",
              "      <th>expected_extracted_people</th>\n",
              "      <th>rationale</th>\n",
              "      <th>extracted_people</th>\n",
              "      <th>extraction_correctness_metric</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>[SOCCER, -, JAPAN, GET, LUCKY, WIN, ,, CHINA, IN, SURPRISE, DEFEAT...</td>\n",
              "      <td>[CHINA]</td>\n",
              "      <td>We extracted \"JAPAN\" and \"CHINA\" as they refer to specific countri...</td>\n",
              "      <td>[JAPAN, CHINA]</td>\n",
              "      <td></td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>[Nadim, Ladki]</td>\n",
              "      <td>[Nadim, Ladki]</td>\n",
              "      <td>We extracted the tokens \"Nadim\" and \"Ladki\" as they refer to speci...</td>\n",
              "      <td>[Nadim, Ladki]</td>\n",
              "      <td>✔️ [True]</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>[AL-AIN, ,, United, Arab, Emirates, 1996-12-06]</td>\n",
              "      <td>[]</td>\n",
              "      <td>There are no tokens referring to specific people in the provided l...</td>\n",
              "      <td>[]</td>\n",
              "      <td>✔️ [True]</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>[Japan, began, the, defence, of, their, Asian, Cup, title, with, a...</td>\n",
              "      <td>[]</td>\n",
              "      <td>We did not find any tokens referring to specific people in the pro...</td>\n",
              "      <td>[]</td>\n",
              "      <td>✔️ [True]</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>[But, China, saw, their, luck, desert, them, in, the, second, matc...</td>\n",
              "      <td>[]</td>\n",
              "      <td>The extracted tokens referring to specific people are \"China\" and ...</td>\n",
              "      <td>[China, Uzbekistan]</td>\n",
              "      <td></td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>...</th>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>195</th>\n",
              "      <td>['The', 'Wallabies', 'have', 'their', 'sights', 'set', 'on', 'a', ...</td>\n",
              "      <td>[David, Campese]</td>\n",
              "      <td>The extracted_people includes \"David Campese\" as it refers to a sp...</td>\n",
              "      <td>[David, Campese]</td>\n",
              "      <td>✔️ [True]</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>196</th>\n",
              "      <td>['The', 'Wallabies', 'currently', 'have', 'no', 'plans', 'to', 'ma...</td>\n",
              "      <td>[]</td>\n",
              "      <td>The extracted_people includes \"Wallabies\" as it refers to a specif...</td>\n",
              "      <td>[]</td>\n",
              "      <td>✔️ [True]</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>197</th>\n",
              "      <td>['Campese', 'will', 'be', 'up', 'against', 'a', 'familiar', 'foe',...</td>\n",
              "      <td>[Campese, Rob, Andrew]</td>\n",
              "      <td>The extracted tokens refer to specific people mentioned in the tex...</td>\n",
              "      <td>[Campese, Rob, Andrew]</td>\n",
              "      <td>✔️ [True]</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>198</th>\n",
              "      <td>['\"', 'Campo', 'has', 'a', 'massive', 'following', 'in', 'this', '...</td>\n",
              "      <td>[Campo, Andrew]</td>\n",
              "      <td>The extracted tokens referring to specific people include \"Campo\" ...</td>\n",
              "      <td>[Campo, Andrew]</td>\n",
              "      <td>✔️ [True]</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>199</th>\n",
              "      <td>['On', 'tour', ',', 'Australia', 'have', 'won', 'all', 'four', 'te...</td>\n",
              "      <td>[]</td>\n",
              "      <td>We extracted the names of specific people from the tokenized text....</td>\n",
              "      <td>[]</td>\n",
              "      <td>✔️ [True]</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "<p>200 rows × 5 columns</p>\n",
              "</div>"
            ],
            "text/plain": [
              "                                                                    tokens  \\\n",
              "0    [SOCCER, -, JAPAN, GET, LUCKY, WIN, ,, CHINA, IN, SURPRISE, DEFEAT...   \n",
              "1                                                           [Nadim, Ladki]   \n",
              "2                          [AL-AIN, ,, United, Arab, Emirates, 1996-12-06]   \n",
              "3    [Japan, began, the, defence, of, their, Asian, Cup, title, with, a...   \n",
              "4    [But, China, saw, their, luck, desert, them, in, the, second, matc...   \n",
              "..                                                                     ...   \n",
              "195  ['The', 'Wallabies', 'have', 'their', 'sights', 'set', 'on', 'a', ...   \n",
              "196  ['The', 'Wallabies', 'currently', 'have', 'no', 'plans', 'to', 'ma...   \n",
              "197  ['Campese', 'will', 'be', 'up', 'against', 'a', 'familiar', 'foe',...   \n",
              "198  ['\"', 'Campo', 'has', 'a', 'massive', 'following', 'in', 'this', '...   \n",
              "199  ['On', 'tour', ',', 'Australia', 'have', 'won', 'all', 'four', 'te...   \n",
              "\n",
              "    expected_extracted_people  \\\n",
              "0                     [CHINA]   \n",
              "1              [Nadim, Ladki]   \n",
              "2                          []   \n",
              "3                          []   \n",
              "4                          []   \n",
              "..                        ...   \n",
              "195          [David, Campese]   \n",
              "196                        []   \n",
              "197    [Campese, Rob, Andrew]   \n",
              "198           [Campo, Andrew]   \n",
              "199                        []   \n",
              "\n",
              "                                                                 rationale  \\\n",
              "0    We extracted \"JAPAN\" and \"CHINA\" as they refer to specific countri...   \n",
              "1    We extracted the tokens \"Nadim\" and \"Ladki\" as they refer to speci...   \n",
              "2    There are no tokens referring to specific people in the provided l...   \n",
              "3    We did not find any tokens referring to specific people in the pro...   \n",
              "4    The extracted tokens referring to specific people are \"China\" and ...   \n",
              "..                                                                     ...   \n",
              "195  The extracted_people includes \"David Campese\" as it refers to a sp...   \n",
              "196  The extracted_people includes \"Wallabies\" as it refers to a specif...   \n",
              "197  The extracted tokens refer to specific people mentioned in the tex...   \n",
              "198  The extracted tokens referring to specific people include \"Campo\" ...   \n",
              "199  We extracted the names of specific people from the tokenized text....   \n",
              "\n",
              "           extracted_people extraction_correctness_metric  \n",
              "0            [JAPAN, CHINA]                                \n",
              "1            [Nadim, Ladki]                     ✔️ [True]  \n",
              "2                        []                     ✔️ [True]  \n",
              "3                        []                     ✔️ [True]  \n",
              "4       [China, Uzbekistan]                                \n",
              "..                      ...                           ...  \n",
              "195        [David, Campese]                     ✔️ [True]  \n",
              "196                      []                     ✔️ [True]  \n",
              "197  [Campese, Rob, Andrew]                     ✔️ [True]  \n",
              "198         [Campo, Andrew]                     ✔️ [True]  \n",
              "199                      []                     ✔️ [True]  \n",
              "\n",
              "[200 rows x 5 columns]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/plain": [
              "86.0"
            ]
          },
          "execution_count": 6,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "evaluate_correctness(people_extractor, devset=test_set)"
      ]
    },
    {
      "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=\"extractor_evaluation\"):\n",
        "    evaluate_correctness = dspy.Evaluate(\n",
        "        devset=test_set,\n",
        "        metric=extraction_correctness_metric,\n",
        "        num_threads=24,\n",
        "        display_progress=True,\n",
        "    )\n",
        "\n",
        "    # Evaluate the program as usual\n",
        "    result = evaluate_correctness(people_extractor)\n",
        "\n",
        "    # Log the aggregated score\n",
        "    mlflow.log_metric(\"exact_match\", result.score)\n",
        "    # Log the detailed evaluation results as a table\n",
        "    mlflow.log_table(\n",
        "        {\n",
        "            \"Tokens\": [example.tokens for example in test_set],\n",
        "            \"Expected\": [example.expected_extracted_people for example in test_set],\n",
        "            \"Predicted\": [output[1] for output in result.results],\n",
        "            \"Exact match\": [output[2] for output in result.results],\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": "markdown",
      "metadata": {},
      "source": [
        "## Optimize the Model\n",
        "\n",
        "DSPy includes powerful optimizers that can improve the quality of your system.\n",
        "\n",
        "Here, we use DSPy's `MIPROv2` optimizer to:\n",
        "- Automatically tune the program's language model (LM) prompt by 1. using the LM to adjust the prompt's instructions and 2. building few-shot examples from the training dataset that are augmented with reasoning generated from `dspy.ChainOfThought`.\n",
        "- Maximize correctness on the training set.\n",
        "\n",
        "This optimization process is automated, saving time and effort while improving accuracy."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "mipro_optimizer = dspy.MIPROv2(\n",
        "    metric=extraction_correctness_metric,\n",
        "    auto=\"medium\",\n",
        ")\n",
        "optimized_people_extractor = mipro_optimizer.compile(\n",
        "    people_extractor,\n",
        "    trainset=train_set,\n",
        "    max_bootstrapped_demos=4,\n",
        "    minibatch=False\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Evaluate Optimized Program\n",
        "\n",
        "After optimization, we re-evaluate the program on the test set to measure improvements. Comparing the optimized and initial results allows us to:\n",
        "- Quantify the benefits of optimization.\n",
        "- Validate that the program generalizes well to unseen data.\n",
        "\n",
        "In this case, we see that accuracy of the program on the test dataset has improved significantly."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Average Metric: 186.00 / 200 (93.0%): 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:23<00:00,  8.58it/s]"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2024/11/18 21:15:00 INFO dspy.evaluate.evaluate: Average Metric: 186 / 200 (93.0%)\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>tokens</th>\n",
              "      <th>expected_extracted_people</th>\n",
              "      <th>rationale</th>\n",
              "      <th>extracted_people</th>\n",
              "      <th>extraction_correctness_metric</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>[SOCCER, -, JAPAN, GET, LUCKY, WIN, ,, CHINA, IN, SURPRISE, DEFEAT...</td>\n",
              "      <td>[CHINA]</td>\n",
              "      <td>There are no specific people mentioned in the provided tokens. The...</td>\n",
              "      <td>[]</td>\n",
              "      <td></td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>[Nadim, Ladki]</td>\n",
              "      <td>[Nadim, Ladki]</td>\n",
              "      <td>The tokens \"Nadim Ladki\" refer to a specific individual. Both toke...</td>\n",
              "      <td>[Nadim, Ladki]</td>\n",
              "      <td>✔️ [True]</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>[AL-AIN, ,, United, Arab, Emirates, 1996-12-06]</td>\n",
              "      <td>[]</td>\n",
              "      <td>There are no tokens referring to specific people in the provided l...</td>\n",
              "      <td>[]</td>\n",
              "      <td>✔️ [True]</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>[Japan, began, the, defence, of, their, Asian, Cup, title, with, a...</td>\n",
              "      <td>[]</td>\n",
              "      <td>There are no specific people mentioned in the provided tokens. The...</td>\n",
              "      <td>[]</td>\n",
              "      <td>✔️ [True]</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>[But, China, saw, their, luck, desert, them, in, the, second, matc...</td>\n",
              "      <td>[]</td>\n",
              "      <td>There are no tokens referring to specific people in the provided l...</td>\n",
              "      <td>[]</td>\n",
              "      <td>✔️ [True]</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>...</th>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>195</th>\n",
              "      <td>['The', 'Wallabies', 'have', 'their', 'sights', 'set', 'on', 'a', ...</td>\n",
              "      <td>[David, Campese]</td>\n",
              "      <td>The extracted tokens refer to a specific person mentioned in the t...</td>\n",
              "      <td>[David, Campese]</td>\n",
              "      <td>✔️ [True]</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>196</th>\n",
              "      <td>['The', 'Wallabies', 'currently', 'have', 'no', 'plans', 'to', 'ma...</td>\n",
              "      <td>[]</td>\n",
              "      <td>There are no specific individuals mentioned in the provided tokens...</td>\n",
              "      <td>[]</td>\n",
              "      <td>✔️ [True]</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>197</th>\n",
              "      <td>['Campese', 'will', 'be', 'up', 'against', 'a', 'familiar', 'foe',...</td>\n",
              "      <td>[Campese, Rob, Andrew]</td>\n",
              "      <td>The tokens include the names \"Campese\" and \"Rob Andrew,\" both of w...</td>\n",
              "      <td>[Campese, Rob, Andrew]</td>\n",
              "      <td>✔️ [True]</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>198</th>\n",
              "      <td>['\"', 'Campo', 'has', 'a', 'massive', 'following', 'in', 'this', '...</td>\n",
              "      <td>[Campo, Andrew]</td>\n",
              "      <td>The extracted tokens refer to specific people mentioned in the tex...</td>\n",
              "      <td>[Campo, Andrew]</td>\n",
              "      <td>✔️ [True]</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>199</th>\n",
              "      <td>['On', 'tour', ',', 'Australia', 'have', 'won', 'all', 'four', 'te...</td>\n",
              "      <td>[]</td>\n",
              "      <td>There are no specific people mentioned in the provided tokens. The...</td>\n",
              "      <td>[]</td>\n",
              "      <td>✔️ [True]</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "<p>200 rows × 5 columns</p>\n",
              "</div>"
            ],
            "text/plain": [
              "                                                                    tokens  \\\n",
              "0    [SOCCER, -, JAPAN, GET, LUCKY, WIN, ,, CHINA, IN, SURPRISE, DEFEAT...   \n",
              "1                                                           [Nadim, Ladki]   \n",
              "2                          [AL-AIN, ,, United, Arab, Emirates, 1996-12-06]   \n",
              "3    [Japan, began, the, defence, of, their, Asian, Cup, title, with, a...   \n",
              "4    [But, China, saw, their, luck, desert, them, in, the, second, matc...   \n",
              "..                                                                     ...   \n",
              "195  ['The', 'Wallabies', 'have', 'their', 'sights', 'set', 'on', 'a', ...   \n",
              "196  ['The', 'Wallabies', 'currently', 'have', 'no', 'plans', 'to', 'ma...   \n",
              "197  ['Campese', 'will', 'be', 'up', 'against', 'a', 'familiar', 'foe',...   \n",
              "198  ['\"', 'Campo', 'has', 'a', 'massive', 'following', 'in', 'this', '...   \n",
              "199  ['On', 'tour', ',', 'Australia', 'have', 'won', 'all', 'four', 'te...   \n",
              "\n",
              "    expected_extracted_people  \\\n",
              "0                     [CHINA]   \n",
              "1              [Nadim, Ladki]   \n",
              "2                          []   \n",
              "3                          []   \n",
              "4                          []   \n",
              "..                        ...   \n",
              "195          [David, Campese]   \n",
              "196                        []   \n",
              "197    [Campese, Rob, Andrew]   \n",
              "198           [Campo, Andrew]   \n",
              "199                        []   \n",
              "\n",
              "                                                                 rationale  \\\n",
              "0    There are no specific people mentioned in the provided tokens. The...   \n",
              "1    The tokens \"Nadim Ladki\" refer to a specific individual. Both toke...   \n",
              "2    There are no tokens referring to specific people in the provided l...   \n",
              "3    There are no specific people mentioned in the provided tokens. The...   \n",
              "4    There are no tokens referring to specific people in the provided l...   \n",
              "..                                                                     ...   \n",
              "195  The extracted tokens refer to a specific person mentioned in the t...   \n",
              "196  There are no specific individuals mentioned in the provided tokens...   \n",
              "197  The tokens include the names \"Campese\" and \"Rob Andrew,\" both of w...   \n",
              "198  The extracted tokens refer to specific people mentioned in the tex...   \n",
              "199  There are no specific people mentioned in the provided tokens. The...   \n",
              "\n",
              "           extracted_people extraction_correctness_metric  \n",
              "0                        []                                \n",
              "1            [Nadim, Ladki]                     ✔️ [True]  \n",
              "2                        []                     ✔️ [True]  \n",
              "3                        []                     ✔️ [True]  \n",
              "4                        []                     ✔️ [True]  \n",
              "..                      ...                           ...  \n",
              "195        [David, Campese]                     ✔️ [True]  \n",
              "196                      []                     ✔️ [True]  \n",
              "197  [Campese, Rob, Andrew]                     ✔️ [True]  \n",
              "198         [Campo, Andrew]                     ✔️ [True]  \n",
              "199                      []                     ✔️ [True]  \n",
              "\n",
              "[200 rows x 5 columns]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/plain": [
              "93.0"
            ]
          },
          "execution_count": 8,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "evaluate_correctness(optimized_people_extractor, devset=test_set)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Inspect Optimized Program's Prompt\n",
        "\n",
        "After optimizing the program, we can inspect the history of interactions to see how DSPy has augmented the program's prompt with few-shot examples. This step demonstrates:\n",
        "- The structure of the prompt used by the program.\n",
        "- How few-shot examples are added to guide the model's behavior.\n",
        "\n",
        "Use `inspect_history(n=1)` to view the last interaction and analyze the generated prompt."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "\n",
            "\n",
            "\n",
            "\u001b[34m[2024-11-18T21:15:00.584497]\u001b[0m\n",
            "\n",
            "\u001b[31mSystem message:\u001b[0m\n",
            "\n",
            "Your input fields are:\n",
            "1. `tokens` (list[str]): tokenized text\n",
            "\n",
            "Your output fields are:\n",
            "1. `rationale` (str): ${produce the extracted_people}. We ...\n",
            "2. `extracted_people` (list[str]): all tokens referring to specific people extracted from the tokenized text\n",
            "\n",
            "All interactions will be structured in the following way, with the appropriate values filled in.\n",
            "\n",
            "[[ ## tokens ## ]]\n",
            "{tokens}\n",
            "\n",
            "[[ ## rationale ## ]]\n",
            "{rationale}\n",
            "\n",
            "[[ ## extracted_people ## ]]\n",
            "{extracted_people}        # note: the value you produce must be pareseable according to the following JSON schema: {\"type\": \"array\", \"items\": {\"type\": \"string\"}}\n",
            "\n",
            "[[ ## completed ## ]]\n",
            "\n",
            "In adhering to this structure, your objective is: \n",
            "        In a high-stakes situation where accurate identification of individuals is critical for regulatory compliance and public health communication, extract contiguous tokens referring to specific people from the provided list of string tokens. Ensure that you output each identified individual as separate tokens without combining multiple tokens into a single value. This task is essential for ensuring clarity and accountability in communications pertaining to EU regulations and health matters.\n",
            "\n",
            "\n",
            "\u001b[31mUser message:\u001b[0m\n",
            "\n",
            "[[ ## tokens ## ]]\n",
            "[\"He\", \"said\", \"further\", \"scientific\", \"study\", \"was\", \"required\", \"and\", \"if\", \"it\", \"was\", \"found\", \"that\", \"action\", \"was\", \"needed\", \"it\", \"should\", \"be\", \"taken\", \"by\", \"the\", \"European\", \"Union\", \".\"]\n",
            "\n",
            "Respond with the corresponding output fields, starting with the field `[[ ## rationale ## ]]`, then `[[ ## extracted_people ## ]]` (must be formatted as a valid Python list[str]), and then ending with the marker for `[[ ## completed ## ]]`.\n",
            "\n",
            "\n",
            "\u001b[31mAssistant message:\u001b[0m\n",
            "\n",
            "[[ ## rationale ## ]]\n",
            "There are no specific people mentioned in the provided tokens. The only reference is to the \"European Union,\" which is an organization rather than a specific person. We therefore conclude that there are no extracted people.\n",
            "\n",
            "[[ ## extracted_people ## ]]\n",
            "[]\n",
            "\n",
            "[[ ## completed ## ]]\n",
            "\n",
            "\n",
            "\u001b[31mUser message:\u001b[0m\n",
            "\n",
            "[[ ## tokens ## ]]\n",
            "[\"BRUSSELS\", \"1996-08-22\"]\n",
            "\n",
            "Respond with the corresponding output fields, starting with the field `[[ ## rationale ## ]]`, then `[[ ## extracted_people ## ]]` (must be formatted as a valid Python list[str]), and then ending with the marker for `[[ ## completed ## ]]`.\n",
            "\n",
            "\n",
            "\u001b[31mAssistant message:\u001b[0m\n",
            "\n",
            "[[ ## rationale ## ]]\n",
            "There are no tokens referring to specific people in the provided list. We only have a location and a date.\n",
            "\n",
            "[[ ## extracted_people ## ]]\n",
            "[]\n",
            "\n",
            "[[ ## completed ## ]]\n",
            "\n",
            "\n",
            "\u001b[31mUser message:\u001b[0m\n",
            "\n",
            "[[ ## tokens ## ]]\n",
            "[\"Fischler\", \"proposed\", \"EU-wide\", \"measures\", \"after\", \"reports\", \"from\", \"Britain\", \"and\", \"France\", \"that\", \"under\", \"laboratory\", \"conditions\", \"sheep\", \"could\", \"contract\", \"Bovine\", \"Spongiform\", \"Encephalopathy\", \"(\", \"BSE\", \")\", \"--\", \"mad\", \"cow\", \"disease\", \".\"]\n",
            "\n",
            "Respond with the corresponding output fields, starting with the field `[[ ## rationale ## ]]`, then `[[ ## extracted_people ## ]]` (must be formatted as a valid Python list[str]), and then ending with the marker for `[[ ## completed ## ]]`.\n",
            "\n",
            "\n",
            "\u001b[31mAssistant message:\u001b[0m\n",
            "\n",
            "[[ ## rationale ## ]]\n",
            "The extracted tokens refer to specific people mentioned in the text. In this case, \"Fischler\" is the only token that refers to a specific person. We extract this token as it is contiguous and represents an individual involved in the proposed measures.\n",
            "\n",
            "[[ ## extracted_people ## ]]\n",
            "[\"Fischler\"]\n",
            "\n",
            "[[ ## completed ## ]]\n",
            "\n",
            "\n",
            "\u001b[31mUser message:\u001b[0m\n",
            "\n",
            "[[ ## tokens ## ]]\n",
            "[\"Germany\", \"'s\", \"representative\", \"to\", \"the\", \"European\", \"Union\", \"'s\", \"veterinary\", \"committee\", \"Werner\", \"Zwingmann\", \"said\", \"on\", \"Wednesday\", \"consumers\", \"should\", \"buy\", \"sheepmeat\", \"from\", \"countries\", \"other\", \"than\", \"Britain\", \"until\", \"the\", \"scientific\", \"advice\", \"was\", \"clearer\", \".\"]\n",
            "\n",
            "Respond with the corresponding output fields, starting with the field `[[ ## rationale ## ]]`, then `[[ ## extracted_people ## ]]` (must be formatted as a valid Python list[str]), and then ending with the marker for `[[ ## completed ## ]]`.\n",
            "\n",
            "\n",
            "\u001b[31mAssistant message:\u001b[0m\n",
            "\n",
            "[[ ## rationale ## ]]\n",
            "We extracted the name \"Werner Zwingmann\" from the tokenized text as it refers to a specific person mentioned in the context of a statement regarding sheepmeat consumption.\n",
            "\n",
            "[[ ## extracted_people ## ]]\n",
            "[\"Werner\", \"Zwingmann\"]\n",
            "\n",
            "[[ ## completed ## ]]\n",
            "\n",
            "\n",
            "\u001b[31mUser message:\u001b[0m\n",
            "\n",
            "[[ ## tokens ## ]]\n",
            "[\"LONDON\", \"1996-12-06\"]\n",
            "\n",
            "Respond with the corresponding output fields, starting with the field `[[ ## rationale ## ]]`, then `[[ ## extracted_people ## ]]` (must be formatted as a valid Python list[str]), and then ending with the marker for `[[ ## completed ## ]]`.\n",
            "\n",
            "\n",
            "\u001b[31mResponse:\u001b[0m\n",
            "\n",
            "\u001b[32m[[ ## rationale ## ]]\n",
            "There are no tokens referring to specific people in the provided list. The tokens only include a location and a date.\n",
            "\n",
            "[[ ## extracted_people ## ]]\n",
            "[]\n",
            "\n",
            "[[ ## completed ## ]]\u001b[0m\n",
            "\n",
            "\n",
            "\n",
            "\n",
            "\n"
          ]
        }
      ],
      "source": [
        "dspy.inspect_history(n=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Keeping an eye on cost\n",
        "\n",
        "DSPy allows you to track the cost of your programs. The following code demonstrates how to obtain the cost of all LM calls made by the DSPy extractor program so far."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "0.26362742999999983"
            ]
          },
          "execution_count": 10,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "cost = sum([x['cost'] for x in lm.history if x['cost'] is not None])  # cost in USD, as calculated by LiteLLM for certain providers\n",
        "cost"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Saving and Loading Optimized Programs\n",
        "\n",
        "DSPy supports saving and loading programs, enabling you to reuse optimized systems without the need to re-optimize from scratch. This feature is especially useful for deploying your programs in production environments or sharing them with collaborators.\n",
        "\n",
        "In this step, we'll save the optimized program to a file and demonstrate how to load it back for future use."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "['Marcello', 'Cuttitta']"
            ]
          },
          "execution_count": 11,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "optimized_people_extractor.save(\"optimized_extractor.json\")\n",
        "\n",
        "loaded_people_extractor = dspy.ChainOfThought(PeopleExtraction)\n",
        "loaded_people_extractor.load(\"optimized_extractor.json\")\n",
        "\n",
        "loaded_people_extractor(tokens=[\"Italy\", \"recalled\", \"Marcello\", \"Cuttitta\"]).extracted_people"
      ]
    },
    {
      "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_extractor\"):\n",
        "    model_info = mlflow.dspy.log_model(\n",
        "        optimized_people_extractor,\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": [
        "## Conclusion\n",
        "\n",
        "In this tutorial, we demonstrated how to:\n",
        "- Use DSPy to build a modular, interpretable system for entity extraction.\n",
        "- Evaluate and optimize the system using DSPy's built-in tools.\n",
        "\n",
        "By leveraging structured inputs and outputs, we ensured that the system was easy to understand and improve. The optimization process allowed us to quickly improve performance without manually crafting prompts or tweaking parameters.\n",
        "\n",
        "**Next Steps:**\n",
        "- Experiment with extraction of other entity types (e.g., locations or organizations).\n",
        "- Explore DSPy's other builtin modules like `ReAct` for more complex reasoning tasks.\n",
        "- Use the system in larger workflows, such as large scale document processing or summarization."
      ]
    }
  ],
  "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.10.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 4
}
