DSPy-Text2SQL.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dspy\n",
    "import random\n",
    "from dotenv import load_dotenv\n",
    "from dspy.datasets import DataLoader\n",
    "from dspy.evaluate import Evaluate\n",
    "from dspy.teleprompt import BootstrapFewShotWithRandomSearch, LabeledFewShot\n",
    "\n",
    "from src.starling import StarlingLM # <- Custom Local Model Client for Starling7B\n",
    "\n",
    "_ = load_dotenv()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### LLM"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model that will be used in this notebook is [Starling7B](https://huggingface.co/Nexusflow/Starling-LM-7B-beta), and the evaluation model will be [GPT-4 Turbo](https://openai.com/gpt-4)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Share generation args between models\n",
    "generation_args = {\n",
    "    \"temperature\":0,\n",
    "    \"max_tokens\":500,\n",
    "    \"stop\":\"\\n\\n\",\n",
    "    \"model_type\":\"chat\",\n",
    "    \"n\": 1\n",
    "}\n",
    "# Model specific args\n",
    "model_info = {\n",
    "    \"gpt-4\": {\"model\": \"gpt-4-0125-preview\", \"api_base\": \"https://api.openai.com/v1/\"},\n",
    "    \"starling\": {\"model\": \"Nexusflow/Starling-LM-7B-beta\"}\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up the models\n",
    "lm = StarlingLM(**model_info[\"starling\"], **generation_args)\n",
    "evaluator_lm = dspy.OpenAI(**model_info[\"gpt-4\"], **generation_args)\n",
    "\n",
    "dspy.configure(lm=lm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[' The capital of Colombia is Bogotá. It is the largest city in the country and serves as the political, economic, and cultural center of Colombia. Located in the Andean region of the country, Bogotá has a rich history and is known for its vibrant arts scene, diverse architecture, and numerous museums and cultural institutions.']"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Testing inference of Starling\n",
    "lm(\"What is the capital of Colombia?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The dataset that will be used in this notebook is [gretelai/synthetic_text_to_sql](https://huggingface.co/datasets/gretelai/synthetic_text_to_sql)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define random seed\n",
    "random.seed(1399)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(75, 25, 75)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load dataset\n",
    "dl = DataLoader()\n",
    "trainset = dl.from_huggingface(\n",
    "    dataset_name=\"gretelai/synthetic_text_to_sql\", # Dataset name from Huggingface\n",
    "    fields=(\"sql_prompt\", \"sql_context\", \"sql\"), # Fields needed\n",
    "    input_keys=(\"sql_prompt\", \"sql_context\"), # What our model expects to recieve to generate an output\n",
    "    split=\"train\"\n",
    ")\n",
    "\n",
    "testset = dl.from_huggingface(\n",
    "    dataset_name=\"gretelai/synthetic_text_to_sql\", # Dataset name from Huggingface\n",
    "    fields=(\"sql_prompt\", \"sql_context\", \"sql\"), # Fields needed\n",
    "    input_keys=(\"sql_prompt\", \"sql_context\"), # What our model expects to recieve to generate an output\n",
    "    split=\"test\"\n",
    ")\n",
    "\n",
    "trainset = dl.sample(dataset=trainset, n=100)\n",
    "testset = dl.sample(dataset=testset, n=75)\n",
    "\n",
    "_trainval = dl.train_test_split(dataset=trainset, test_size=0.25, random_state=1399) # 25% of training data for validation\n",
    "trainset, valset = _trainval[\"train\"], _trainval[\"test\"]\n",
    "\n",
    "len(trainset), len(valset), len(testset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "SQL_PROMPT:\n",
      "\n",
      "List the top 5 countries with the highest number of satellites in orbit as of 2022-01-01, ordered by the number of satellites in descending order.\n",
      "\n",
      "SQL_CONTEXT:\n",
      "\n",
      "CREATE TABLE countries(id INT, name VARCHAR(255), population INT, satellites_in_orbit INT, last_census_date DATE);\n",
      "\n",
      "SQL:\n",
      "\n",
      "SELECT name, satellites_in_orbit FROM countries WHERE last_census_date <= '2022-01-01' GROUP BY name ORDER BY satellites_in_orbit DESC LIMIT 5;\n"
     ]
    }
   ],
   "source": [
    "# Verify an example of the dataset\n",
    "sample = dl.sample(dataset=trainset, n=1)[0]\n",
    "for k, v in sample.items():\n",
    "    print(f\"\\n{k.upper()}:\\n\")\n",
    "    print(v)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Signature (Input/Output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TextToSql(dspy.Signature):\n",
    "    \"\"\"Transform a natural language query into a SQL query.\"\"\"\n",
    "\n",
    "    sql_prompt = dspy.InputField(desc=\"Natural language query\")\n",
    "    sql_context = dspy.InputField(desc=\"Context for the query\")\n",
    "    sql = dspy.OutputField(desc=\"SQL query\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Baseline Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "SQL:\n",
      "\n",
      "SELECT name, satellites_in_orbit\n",
      "FROM countries\n",
      "WHERE last_census_date = '2022-01-01'\n",
      "ORDER BY satellites_in_orbit DESC\n",
      "LIMIT 5;\n"
     ]
    }
   ],
   "source": [
    "generate_sql_query = dspy.Predict(signature=TextToSql)\n",
    "\n",
    "result = generate_sql_query(\n",
    "    sql_prompt=sample[\"sql_prompt\"],\n",
    "    sql_context=sample[\"sql_context\"]\n",
    ")\n",
    "\n",
    "for k, v in result.items():\n",
    "    print(f\"\\n{k.upper()}:\\n\")\n",
    "    print(v)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### ChainOfThought Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "RATIONALE:\n",
      "\n",
      "produce the SQL query. We need to:\n",
      "\n",
      "SQL:\n",
      "\n",
      "1. Select the relevant columns from the table: In this case, we need to select the country name and the number of satellites in orbit.\n",
      "2. Filter the data based on the date: We need to filter the data to only include records with a last_census_date of 2022-01-01.\n",
      "3. Order the results: We need to order the results by the number of satellites in orbit in descending order.\n"
     ]
    }
   ],
   "source": [
    "generate_sql_query = dspy.ChainOfThought(signature=TextToSql)\n",
    "\n",
    "result = generate_sql_query(\n",
    "    sql_prompt=sample[\"sql_prompt\"],\n",
    "    sql_context=sample[\"sql_context\"]\n",
    ")\n",
    "\n",
    "for k, v in result.items():\n",
    "    print(f\"\\n{k.upper()}:\\n\")\n",
    "    print(v)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Metric of evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Metric definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Correctness(dspy.Signature):\n",
    "    \"\"\"Assess if the SQL query accurately answers the given natural language query based on the provided context.\"\"\"\n",
    "\n",
    "    sql_prompt = dspy.InputField(desc=\"Natural language query \")\n",
    "    sql_context = dspy.InputField(desc=\"Context for the query\")\n",
    "    sql = dspy.InputField(desc=\"SQL query\")\n",
    "    correct = dspy.OutputField(desc=\"Indicate whether the SQL query correctly answers the natural language query based on the given context\", prefix=\"Yes/No:\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def correctness_metric(example, pred, trace=None):\n",
    "    sql_prompt, sql_context, sql = example.sql_prompt, example.sql_context, pred.sql\n",
    "\n",
    "    correctness = dspy.Predict(Correctness)\n",
    "\n",
    "    with dspy.context(lm=evaluator_lm): \n",
    "        correct = correctness(\n",
    "            sql_prompt=sql_prompt,\n",
    "            sql_context=sql_context,\n",
    "            sql=sql,\n",
    "        )\n",
    "    \n",
    "    score = int(correct.correct==\"Yes\")\n",
    "\n",
    "    if trace is not None:\n",
    "        return score == 1\n",
    "\n",
    "    return score"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Evaluate single data point"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Correct SQL query: No\n"
     ]
    }
   ],
   "source": [
    "_correctness = correctness_metric(\n",
    "    example=sample,\n",
    "    pred=result\n",
    ")\n",
    "print(f\"Correct SQL query: {'Yes' if _correctness else 'No'}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\n",
      "\n",
      "Assess if the SQL query accurately answers the given natural language query based on the provided context.\n",
      "\n",
      "---\n",
      "\n",
      "Follow the following format.\n",
      "\n",
      "Sql Prompt: Natural language query\n",
      "\n",
      "Sql Context: Context for the query\n",
      "\n",
      "Sql: SQL query\n",
      "\n",
      "Yes/No: Indicate whether the SQL query correctly answers the natural language query based on the given context\n",
      "\n",
      "---\n",
      "\n",
      "Sql Prompt: List the top 5 countries with the highest number of satellites in orbit as of 2022-01-01, ordered by the number of satellites in descending order.\n",
      "\n",
      "Sql Context: CREATE TABLE countries(id INT, name VARCHAR(255), population INT, satellites_in_orbit INT, last_census_date DATE);\n",
      "\n",
      "Sql: 1. Select the relevant columns from the table: In this case, we need to select the country name and the number of satellites in orbit. 2. Filter the data based on the date: We need to filter the data to only include records with a last_census_date of 2022-01-01. 3. Order the results: We need to order the results by the number of satellites in orbit in descending order.\n",
      "\n",
      "Yes/No:\u001b[32m No\u001b[0m\n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "evaluator_lm.inspect_history(n=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Evaluate entire dataset - GPT 3.5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #F0F0F0; padding: 10px; border-radius: 5px;\"> <p style=\"color: #4B4B4B; font-size: 18px; font-weight: bold; margin: 0;\"> 📊 Baseline Evaluation </p> <p style=\"color: #4B4B4B; font-size: 16px; margin: 5px 0 0;\"> Without any optimization, <strong>Starling7B</strong> achieves an <strong>80% correctness in validation (25 samples)</strong> and <strong>70.07% correctness in test (75 samples).</strong> </p> </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPT 3.5 Turbo - Validation Score: \n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/25 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Average Metric: 20 / 25  (80.0): 100%|██████████| 25/25 [00:09<00:00,  2.77it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Metric: 20 / 25  (80.0%)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "/home/jjmov99/dspy-testing/.venv/lib/python3.11/site-packages/dspy/evaluate/evaluate.py:187: FutureWarning: DataFrame.applymap has been deprecated. Use DataFrame.map instead.\n",
      "  df = df.applymap(truncate_cell)\n"
     ]
    }
   ],
   "source": [
    "print(\"GPT 3.5 Turbo - Validation Score: \\n\")\n",
    "with dspy.context(lm=dspy.OpenAI(model=\"gpt-3.5-turbo-0125\", api_base=\"https://api.openai.com/v1/\", **generation_args)):\n",
    "    evaluate = Evaluate(devset=valset, metric=correctness_metric, num_threads=10, display_progress=True, display_table=0)\n",
    "    evaluate(generate_sql_query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPT 3.5 Turbo - Test Score: \n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Average Metric: 53 / 75  (70.7): 100%|██████████| 75/75 [00:14<00:00,  5.06it/s] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Metric: 53 / 75  (70.7%)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "/home/jjmov99/dspy-testing/.venv/lib/python3.11/site-packages/dspy/evaluate/evaluate.py:187: FutureWarning: DataFrame.applymap has been deprecated. Use DataFrame.map instead.\n",
      "  df = df.applymap(truncate_cell)\n"
     ]
    }
   ],
   "source": [
    "print(\"GPT 3.5 Turbo - Test Score: \\n\")\n",
    "with dspy.context(lm=dspy.OpenAI(model=\"gpt-3.5-turbo-0125\", api_base=\"https://api.openai.com/v1/\", **generation_args)):\n",
    "    evaluate = Evaluate(devset=testset, metric=correctness_metric, num_threads=10, display_progress=True, display_table=0)\n",
    "    evaluate(generate_sql_query)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Evaluate entire dataset - Starling7B"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #FFCCCB; padding: 10px; border-radius: 5px;\"> <p style=\"color: #8B0000; font-size: 18px; font-weight: bold; margin: 0;\"> ⚠️ Evaluation Stage 1 </p> <p style=\"color: #8B0000; font-size: 16px; margin: 5px 0 0;\"> Without any optimization, <strong>Starling7B</strong> achieves an <strong>72% correctness in validation (25 samples)</strong> and <strong>50.67% correctness in test (75 samples).</strong> </p> </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starling7b - Validation Score: \n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Average Metric: 18 / 25  (72.0): 100%|██████████| 25/25 [00:10<00:00,  2.44it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Metric: 18 / 25  (72.0%)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "/home/jjmov99/dspy-testing/.venv/lib/python3.11/site-packages/dspy/evaluate/evaluate.py:187: FutureWarning: DataFrame.applymap has been deprecated. Use DataFrame.map instead.\n",
      "  df = df.applymap(truncate_cell)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "72.0"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(\"Starling7b - Validation Score: \\n\")\n",
    "evaluate = Evaluate(devset=valset, metric=correctness_metric, num_threads=10, display_progress=True, display_table=0)\n",
    "evaluate(generate_sql_query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starling7b - Test Score: \n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Average Metric: 38 / 75  (50.7): 100%|██████████| 75/75 [00:31<00:00,  2.37it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Metric: 38 / 75  (50.7%)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "/home/jjmov99/dspy-testing/.venv/lib/python3.11/site-packages/dspy/evaluate/evaluate.py:187: FutureWarning: DataFrame.applymap has been deprecated. Use DataFrame.map instead.\n",
      "  df = df.applymap(truncate_cell)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "50.67"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(\"Starling7b - Test Score: \\n\")\n",
    "evaluate = Evaluate(devset=testset, metric=correctness_metric, num_threads=10, display_progress=True, display_table=0)\n",
    "evaluate(generate_sql_query)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimize for Text2SQL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create program"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the program ~ You can think of this a Pytorch model.\n",
    "class TextToSqlProgram(dspy.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.program = dspy.ChainOfThought(signature=TextToSql)\n",
    "    \n",
    "    def forward(self, sql_prompt, sql_context):\n",
    "        return self.program(\n",
    "            sql_prompt=sql_prompt,\n",
    "            sql_context=sql_context\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### FewShot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Execute the optimizer -> this only adds few shots to the prompt\n",
    "optimizer = LabeledFewShot(k=4)\n",
    "optmized_program = optimizer.compile(student=TextToSqlProgram(), trainset=trainset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Prediction(\n",
       "    rationale='produce the SQL query. We need to filter the countries based on the date and then order them by the number of satellites in orbit.',\n",
       "    sql=\"SELECT name, satellites_in_orbit\\nFROM countries\\nWHERE last_census_date <= '2022-01-01'\\nORDER BY satellites_in_orbit DESC\\nLIMIT 5;\"\n",
       ")"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "optmized_program(sql_context=sample[\"sql_context\"], sql_prompt=sample[\"sql_prompt\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### What is happening inside?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\n",
      "\n",
      "Transform a natural language query into a SQL query.\n",
      "\n",
      "---\n",
      "\n",
      "Sql Prompt: What is the total number of electric vehicles sold by manufacturer 'XYZ'?\n",
      "Sql Context: CREATE TABLE sales_data (manufacturer VARCHAR(10), vehicle_type VARCHAR(10), quantity INT);\n",
      "Sql: SELECT manufacturer, SUM(quantity) FROM sales_data WHERE vehicle_type = 'Electric' AND manufacturer = 'XYZ' GROUP BY manufacturer;\n",
      "\n",
      "Sql Prompt: What is the total number of amphibians in the 'animals' table with a population size greater than 1000?\n",
      "Sql Context: CREATE TABLE animals (id INT, name VARCHAR(50), species VARCHAR(50), population_size INT); INSERT INTO animals (id, name, species, population_size) VALUES (1, 'Frog', 'Anura', 1200);\n",
      "Sql: SELECT COUNT(*) FROM animals WHERE species = 'Anura' AND population_size > 1000;\n",
      "\n",
      "Sql Prompt: What is the minimum depth a marine species can live at?\n",
      "Sql Context: CREATE TABLE species (id INT, name VARCHAR(255), habitat VARCHAR(255), depth FLOAT); INSERT INTO species (id, name, habitat, depth) VALUES (1, 'Clownfish', 'Coral Reef', 20.0); INSERT INTO species (id, name, habitat, depth) VALUES (2, 'Blue Whale', 'Open Ocean', 2000.0); INSERT INTO species (id, name, habitat, depth) VALUES (3, 'Sea Otter', 'Kelp Forest', 50.0);\n",
      "Sql: SELECT MIN(depth) FROM species;\n",
      "\n",
      "Sql Prompt: Delete all investments with ESG scores less than 70.\n",
      "Sql Context: CREATE TABLE investments (id INT, sector VARCHAR(20), esg_score FLOAT); INSERT INTO investments (id, sector, esg_score) VALUES (1, 'Education', 75.00), (2, 'Healthcare', 70.00), (3, 'Renewable Energy', 65.00);\n",
      "Sql: DELETE FROM investments WHERE esg_score < 70;\n",
      "\n",
      "---\n",
      "\n",
      "Follow the following format.\n",
      "\n",
      "Sql Prompt: Natural language query\n",
      "\n",
      "Sql Context: Context for the query\n",
      "\n",
      "Reasoning: Let's think step by step in order to ${produce the sql}. We ...\n",
      "\n",
      "Sql: SQL query\n",
      "\n",
      "---\n",
      "\n",
      "Sql Prompt: List the top 5 countries with the highest number of satellites in orbit as of 2022-01-01, ordered by the number of satellites in descending order.\n",
      "\n",
      "Sql Context: CREATE TABLE countries(id INT, name VARCHAR(255), population INT, satellites_in_orbit INT, last_census_date DATE);\n",
      "\n",
      "Reasoning: Let's think step by step in order to produce the SQL query. We need to filter the countries based on the date and then order them by the number of satellites in orbit.\n",
      "\n",
      "Sql:\u001b[32m SELECT name, satellites_in_orbit\n",
      "FROM countries\n",
      "WHERE last_census_date <= '2022-01-01'\n",
      "ORDER BY satellites_in_orbit DESC\n",
      "LIMIT 5;\u001b[0m\n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "lm.inspect_history(n=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Evaluate the optimized program"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "<div style=\"background-color: #FFF8DC; padding: 10px; border-radius: 5px;\"> <p style=\"color: #DAA520; font-size: 18px; font-weight: bold; margin: 0;\"> 🌟 Evaluation Stage 2 </p> <p style=\"color: #DAA520; font-size: 16px; margin: 5px 0 0;\"> With <strong>Few Shot</strong> optimization, <strong>Starling7B</strong> achieves an <strong>64% correctness in validation (25 samples)</strong> and <strong>60% correctness in test (75 samples).</strong> </p> </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starling7b + FewShotOptimizer - Validation Score: \n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Average Metric: 16 / 25  (64.0): 100%|██████████| 25/25 [00:15<00:00,  1.59it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Metric: 16 / 25  (64.0%)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "/home/jjmov99/dspy-testing/.venv/lib/python3.11/site-packages/dspy/evaluate/evaluate.py:187: FutureWarning: DataFrame.applymap has been deprecated. Use DataFrame.map instead.\n",
      "  df = df.applymap(truncate_cell)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "64.0"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(\"Starling7b + FewShotOptimizer - Validation Score: \\n\")\n",
    "evaluate = Evaluate(devset=valset, metric=correctness_metric, num_threads=10, display_progress=True, display_table=0)\n",
    "evaluate(optmized_program)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starling7b + FewShotOptimizer - Test Score: \n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Average Metric: 45 / 75  (60.0): 100%|██████████| 75/75 [00:31<00:00,  2.35it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Metric: 45 / 75  (60.0%)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "60.0"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(\"Starling7b + FewShotOptimizer - Test Score: \\n\")\n",
    "evaluate = Evaluate(devset=testset, metric=correctness_metric, num_threads=10, display_progress=True, display_table=0)\n",
    "evaluate(optmized_program)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### BootstrapFewShotWithRandomSearch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[DSPy docs](https://dspy-docs.vercel.app/docs/building-blocks/optimizers) recommend that in a setup like the one with have at hand, with ~50 samples, the best option is to use `BootstrapFewShotWithRandomSearch`:\n",
    "\n",
    "![image](assets/dspy.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Going to sample between 1 and 2 traces per predictor.\n",
      "Will attempt to train 8 candidate sets.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Average Metric: 18 / 25  (72.0): 100%|██████████| 25/25 [00:15<00:00,  1.64it/s]\n",
      "/home/jjmov99/dspy-testing/.venv/lib/python3.11/site-packages/dspy/evaluate/evaluate.py:187: FutureWarning: DataFrame.applymap has been deprecated. Use DataFrame.map instead.\n",
      "  df = df.applymap(truncate_cell)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Metric: 18 / 25  (72.0%)\n",
      "Score: 72.0 for set: [0]\n",
      "New best score: 72.0 for seed -3\n",
      "Scores so far: [72.0]\n",
      "Best score: 72.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Average Metric: 19 / 25  (76.0): 100%|██████████| 25/25 [00:24<00:00,  1.00it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Metric: 19 / 25  (76.0%)\n",
      "Score: 76.0 for set: [16]\n",
      "New best score: 76.0 for seed -2\n",
      "Scores so far: [72.0, 76.0]\n",
      "Best score: 76.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  4%|▍         | 3/75 [00:11<04:44,  3.95s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Bootstrapped 2 full traces after 4 examples in round 0.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Average Metric: 18 / 25  (72.0): 100%|██████████| 25/25 [00:25<00:00,  1.02s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Metric: 18 / 25  (72.0%)\n",
      "Score: 72.0 for set: [16]\n",
      "Scores so far: [72.0, 76.0, 72.0]\n",
      "Best score: 76.0\n",
      "Average of max per entry across top 1 scores: 0.76\n",
      "Average of max per entry across top 2 scores: 0.84\n",
      "Average of max per entry across top 3 scores: 0.84\n",
      "Average of max per entry across top 5 scores: 0.84\n",
      "Average of max per entry across top 8 scores: 0.84\n",
      "Average of max per entry across top 9999 scores: 0.84\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  4%|▍         | 3/75 [00:10<04:08,  3.45s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Bootstrapped 2 full traces after 4 examples in round 0.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Average Metric: 20 / 25  (80.0): 100%|██████████| 25/25 [00:26<00:00,  1.07s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Metric: 20 / 25  (80.0%)\n",
      "Score: 80.0 for set: [16]\n",
      "New best score: 80.0 for seed 0\n",
      "Scores so far: [72.0, 76.0, 72.0, 80.0]\n",
      "Best score: 80.0\n",
      "Average of max per entry across top 1 scores: 0.8\n",
      "Average of max per entry across top 2 scores: 0.84\n",
      "Average of max per entry across top 3 scores: 0.88\n",
      "Average of max per entry across top 5 scores: 0.88\n",
      "Average of max per entry across top 8 scores: 0.88\n",
      "Average of max per entry across top 9999 scores: 0.88\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  1%|▏         | 1/75 [00:02<03:34,  2.90s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Bootstrapped 1 full traces after 2 examples in round 0.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Average Metric: 20 / 25  (80.0): 100%|██████████| 25/25 [00:25<00:00,  1.01s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Metric: 20 / 25  (80.0%)\n",
      "Score: 80.0 for set: [16]\n",
      "Scores so far: [72.0, 76.0, 72.0, 80.0, 80.0]\n",
      "Best score: 80.0\n",
      "Average of max per entry across top 1 scores: 0.8\n",
      "Average of max per entry across top 2 scores: 0.88\n",
      "Average of max per entry across top 3 scores: 0.88\n",
      "Average of max per entry across top 5 scores: 0.88\n",
      "Average of max per entry across top 8 scores: 0.88\n",
      "Average of max per entry across top 9999 scores: 0.88\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  1%|▏         | 1/75 [00:03<04:09,  3.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Bootstrapped 1 full traces after 2 examples in round 0.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Average Metric: 20 / 25  (80.0): 100%|██████████| 25/25 [00:26<00:00,  1.08s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Metric: 20 / 25  (80.0%)\n",
      "Score: 80.0 for set: [16]\n",
      "Scores so far: [72.0, 76.0, 72.0, 80.0, 80.0, 80.0]\n",
      "Best score: 80.0\n",
      "Average of max per entry across top 1 scores: 0.8\n",
      "Average of max per entry across top 2 scores: 0.88\n",
      "Average of max per entry across top 3 scores: 0.96\n",
      "Average of max per entry across top 5 scores: 0.96\n",
      "Average of max per entry across top 8 scores: 0.96\n",
      "Average of max per entry across top 9999 scores: 0.96\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  1%|▏         | 1/75 [00:03<04:00,  3.25s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Bootstrapped 1 full traces after 2 examples in round 0.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Average Metric: 19 / 25  (76.0): 100%|██████████| 25/25 [00:23<00:00,  1.05it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Metric: 19 / 25  (76.0%)\n",
      "Score: 76.0 for set: [16]\n",
      "Scores so far: [72.0, 76.0, 72.0, 80.0, 80.0, 80.0, 76.0]\n",
      "Best score: 80.0\n",
      "Average of max per entry across top 1 scores: 0.8\n",
      "Average of max per entry across top 2 scores: 0.88\n",
      "Average of max per entry across top 3 scores: 0.96\n",
      "Average of max per entry across top 5 scores: 0.96\n",
      "Average of max per entry across top 8 scores: 0.96\n",
      "Average of max per entry across top 9999 scores: 0.96\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  1%|▏         | 1/75 [00:03<04:18,  3.49s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Bootstrapped 1 full traces after 2 examples in round 0.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Average Metric: 20 / 25  (80.0): 100%|██████████| 25/25 [00:25<00:00,  1.04s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Metric: 20 / 25  (80.0%)\n",
      "Score: 80.0 for set: [16]\n",
      "Scores so far: [72.0, 76.0, 72.0, 80.0, 80.0, 80.0, 76.0, 80.0]\n",
      "Best score: 80.0\n",
      "Average of max per entry across top 1 scores: 0.8\n",
      "Average of max per entry across top 2 scores: 0.88\n",
      "Average of max per entry across top 3 scores: 0.96\n",
      "Average of max per entry across top 5 scores: 1.0\n",
      "Average of max per entry across top 8 scores: 1.0\n",
      "Average of max per entry across top 9999 scores: 1.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  5%|▌         | 4/75 [00:13<03:57,  3.34s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Bootstrapped 2 full traces after 5 examples in round 0.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Average Metric: 20 / 25  (80.0): 100%|██████████| 25/25 [00:25<00:00,  1.02s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Metric: 20 / 25  (80.0%)\n",
      "Score: 80.0 for set: [16]\n",
      "Scores so far: [72.0, 76.0, 72.0, 80.0, 80.0, 80.0, 76.0, 80.0, 80.0]\n",
      "Best score: 80.0\n",
      "Average of max per entry across top 1 scores: 0.8\n",
      "Average of max per entry across top 2 scores: 0.88\n",
      "Average of max per entry across top 3 scores: 0.96\n",
      "Average of max per entry across top 5 scores: 1.0\n",
      "Average of max per entry across top 8 scores: 1.0\n",
      "Average of max per entry across top 9999 scores: 1.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  1%|▏         | 1/75 [00:03<04:52,  3.95s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Bootstrapped 1 full traces after 2 examples in round 0.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Average Metric: 19 / 25  (76.0): 100%|██████████| 25/25 [00:27<00:00,  1.10s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Metric: 19 / 25  (76.0%)\n",
      "Score: 76.0 for set: [16]\n",
      "Scores so far: [72.0, 76.0, 72.0, 80.0, 80.0, 80.0, 76.0, 80.0, 80.0, 76.0]\n",
      "Best score: 80.0\n",
      "Average of max per entry across top 1 scores: 0.8\n",
      "Average of max per entry across top 2 scores: 0.88\n",
      "Average of max per entry across top 3 scores: 0.96\n",
      "Average of max per entry across top 5 scores: 1.0\n",
      "Average of max per entry across top 8 scores: 1.0\n",
      "Average of max per entry across top 9999 scores: 1.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  4%|▍         | 3/75 [00:11<04:41,  3.90s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Bootstrapped 2 full traces after 4 examples in round 0.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Average Metric: 20 / 25  (80.0): 100%|██████████| 25/25 [00:27<00:00,  1.09s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Metric: 20 / 25  (80.0%)\n",
      "Score: 80.0 for set: [16]\n",
      "Scores so far: [72.0, 76.0, 72.0, 80.0, 80.0, 80.0, 76.0, 80.0, 80.0, 76.0, 80.0]\n",
      "Best score: 80.0\n",
      "Average of max per entry across top 1 scores: 0.8\n",
      "Average of max per entry across top 2 scores: 0.88\n",
      "Average of max per entry across top 3 scores: 0.96\n",
      "Average of max per entry across top 5 scores: 1.0\n",
      "Average of max per entry across top 8 scores: 1.0\n",
      "Average of max per entry across top 9999 scores: 1.0\n",
      "11 candidate programs found.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "optimizer2 = BootstrapFewShotWithRandomSearch(metric=correctness_metric, max_bootstrapped_demos=2, num_candidate_programs=8, num_threads=5)\n",
    "optmized_program_2 = optimizer2.compile(student = TextToSqlProgram(), trainset=trainset, valset=valset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Prediction(\n",
       "    rationale='produce the SQL query. We need to find the top 5 countries with the highest number of satellites in orbit as of 2022-01-01, ordered by the number of satellites in descending order.',\n",
       "    sql=\"SELECT name, satellites_in_orbit\\nFROM countries\\nWHERE last_census_date <= '2022-01-01'\\nORDER BY satellites_in_orbit DESC\\nLIMIT 5;\"\n",
       ")"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "optmized_program_2(sql_context=sample[\"sql_context\"], sql_prompt=sample[\"sql_prompt\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\n",
      "\n",
      "Transform a natural language query into a SQL query.\n",
      "\n",
      "---\n",
      "\n",
      "Sql Prompt: What is the total number of open pedagogy courses offered by each country?\n",
      "Sql Context: CREATE TABLE country (country_id INT, country_name VARCHAR(255)); CREATE TABLE open_pedagogy_courses (country_id INT, course_id INT); INSERT INTO country (country_id, country_name) VALUES (6001, 'Country X'), (6002, 'Country Y'), (6003, 'Country Z'); INSERT INTO open_pedagogy_courses (country_id, course_id) VALUES (6001, 7001), (6001, 7002), (6002, 7003);\n",
      "Sql: SELECT country_name, COUNT(course_id) as total_courses FROM country JOIN open_pedagogy_courses ON country.country_id = open_pedagogy_courses.country_id GROUP BY country_name;\n",
      "\n",
      "Sql Prompt: Top 3 most popular songs of 2021 in 'concert_ticket_sales' table?\n",
      "Sql Context: CREATE TABLE concert_ticket_sales (ticket_id INT, song_id INT, quantity INT, price FLOAT, sale_date DATE);\n",
      "Sql: SELECT song_id, SUM(quantity) as total_quantity FROM concert_ticket_sales WHERE sale_date >= '2021-01-01' GROUP BY song_id ORDER BY total_quantity DESC LIMIT 3;\n",
      "\n",
      "Sql Prompt: What is the total amount donated by each donor in the 'Donors' table, ordered by the total amount donated in descending order?\n",
      "Sql Context: CREATE TABLE Donors (DonorID INT, Name VARCHAR(50), TotalDonated DECIMAL(10, 2));\n",
      "Sql: SELECT DonorID, Name, SUM(TotalDonated) AS TotalDonatedSum FROM Donors GROUP BY DonorID ORDER BY TotalDonatedSum DESC;\n",
      "\n",
      "Sql Prompt: What was the total sales amount for each product category in Africa in 2021?\n",
      "Sql Context: CREATE TABLE sales_2021 AS SELECT * FROM sales WHERE sale_date BETWEEN '2021-01-01' AND '2021-12-31'; ALTER TABLE sales_2021 ADD COLUMN sale_country VARCHAR(50); UPDATE sales_2021 SET sale_country = CASE WHEN sale_city IN ('Accra', 'Lagos', 'Cairo') THEN 'Africa' ELSE sale_country END; ALTER TABLE sales_2021 ADD COLUMN product_category VARCHAR(50); UPDATE sales_2021 SET product_category = CASE WHEN product_id = 1 THEN 'Tops' WHEN product_id = 2 THEN 'Bottoms' WHEN product_id = 3 THEN 'Outerwear' WHEN product_id = 4 THEN 'Accessories' END;\n",
      "Sql: SELECT sale_country, product_category, SUM(sale_amount) FROM sales_2021 WHERE sale_country = 'Africa' GROUP BY sale_country, product_category;\n",
      "\n",
      "Sql Prompt: What is the total transaction value for each customer in the past week, split by currency, for customers in the United States?\n",
      "Sql Context: CREATE TABLE customers (customer_id INT, age INT, name VARCHAR(255), country VARCHAR(50)); CREATE TABLE transactions (transaction_id INT, customer_id INT, product_id INT, category_id INT, transaction_date DATE, amount DECIMAL(10,2), currency VARCHAR(10));\n",
      "Sql: SELECT c.country, c.name, t.currency, SUM(t.amount) as total_transaction_value FROM customers c INNER JOIN transactions t ON c.customer_id = t.customer_id WHERE c.country = 'United States' AND t.transaction_date >= DATE_SUB(CURRENT_DATE, INTERVAL 1 WEEK) GROUP BY c.country, c.name, t.currency;\n",
      "\n",
      "Sql Prompt: Who are the top 5 customers by transaction count?\n",
      "Sql Context: CREATE TABLE customer_transactions (customer_id INT, customer_name VARCHAR(20), transaction_id INT); INSERT INTO customer_transactions (customer_id, customer_name, transaction_id) VALUES (1, 'Juan Pérez', 1), (2, 'María Rodríguez', 2), (3, 'Carlos López', 3), (4, 'Laura González', 4), (5, 'José Hernández', 5), (6, 'Ana Sánchez', 6), (7, 'Pedro Martínez', 7);\n",
      "Sql: SELECT c.customer_name, COUNT(ct.transaction_id) as transaction_count FROM customers c JOIN customer_transactions ct ON c.customer_id = ct.customer_id GROUP BY c.customer_name ORDER BY transaction_count DESC LIMIT 5;\n",
      "\n",
      "Sql Prompt: How many accounts are associated with each risk category and what is the total assets value for each category?\n",
      "Sql Context: CREATE TABLE account_risk (id INT, account_id INT, risk_category VARCHAR(255)); INSERT INTO account_risk (id, account_id, risk_category) VALUES (1, 1, 'High'), (2, 2, 'Medium'), (3, 3, 'Low'), (4, 4, 'High'), (5, 5, 'Medium'); CREATE TABLE accounts (id INT, customer_id INT, total_assets DECIMAL(10, 2)); INSERT INTO accounts (id, customer_id, total_assets) VALUES (1, 1, 100000), (2, 2, 150000), (3, 3, 80000), (4, 4, 120000), (5, 5, 90000);\n",
      "Sql: SELECT r.risk_category, COUNT(r.account_id) AS num_accounts, SUM(a.total_assets) AS total_assets FROM account_risk r INNER JOIN accounts a ON r.account_id = a.id GROUP BY r.risk_category;\n",
      "\n",
      "Sql Prompt: What is the average production per acre for corn and soy in each region?\n",
      "Sql Context: CREATE TABLE crop_production (region TEXT, crop_type TEXT, acres FLOAT, production INT); INSERT INTO crop_production (region, crop_type, acres, production) VALUES ('North', 'Corn', 150, 270), ('North', 'Soy', 250, 120), ('South', 'Corn', 220, 900), ('East', 'Soy', 150, 300);\n",
      "Sql: SELECT region, crop_type, AVG(production/acres) as avg_production_per_acre FROM crop_production WHERE crop_type IN ('Corn', 'Soy') GROUP BY region, crop_type;\n",
      "\n",
      "Sql Prompt: What is the minimum policy issuance year for policy number 1001?\n",
      "Sql Context: CREATE TABLE policies (policy_id INT, policy_issue_year INT); INSERT INTO policies (policy_id, policy_issue_year) VALUES (1001, 2018), (1002, 2017), (1003, 2016), (1004, 2019);\n",
      "Sql: SELECT MIN(policy_issue_year) FROM policies WHERE policy_id = 1001;\n",
      "\n",
      "Sql Prompt: What is the total number of hospital beds in rural areas of each state?\n",
      "Sql Context: CREATE TABLE beds (bed_id INT, hospital_id INT, location VARCHAR(20));\n",
      "Sql: SELECT hospital_id, COUNT(*) FROM beds WHERE location = 'Rural' GROUP BY hospital_id;\n",
      "\n",
      "Sql Prompt: What is the average age of employees with the title 'Supervisor' in the 'employees' table?\n",
      "Sql Context: CREATE TABLE employees(id INT, name VARCHAR(255), title VARCHAR(255), age INT); INSERT INTO employees(id, name, title, age) VALUES ('1', 'Jane Smith', 'Mining Supervisor', '55');\n",
      "Sql: SELECT AVG(age) FROM employees WHERE title LIKE '%Supervisor%';\n",
      "\n",
      "Sql Prompt: What was the total number of citizens in Asia in 2017?\n",
      "Sql Context: CREATE TABLE asia_population (id INT PRIMARY KEY, year INT, num_citizens INT); INSERT INTO asia_population (id, year, num_citizens) VALUES (1, 2017, 4000000000);\n",
      "Sql: SELECT num_citizens FROM asia_population WHERE year = 2017;\n",
      "\n",
      "Sql Prompt: What is the average number of citations for algorithmic fairness papers published in 2018 and 2019?\n",
      "Sql Context: CREATE TABLE algorithmic_fairness_papers (year INT, paper_title VARCHAR(255), author_name VARCHAR(255), num_citations INT); INSERT INTO algorithmic_fairness_papers (year, paper_title, author_name, num_citations) VALUES ('2018', 'Algorithmic Fairness: A Review', 'Alice Johnson', '50');\n",
      "Sql: SELECT AVG(num_citations) as avg_citations FROM algorithmic_fairness_papers WHERE year IN (2018, 2019);\n",
      "\n",
      "Sql Prompt: What is the average water temperature in salmon farms in Norway?\n",
      "Sql Context: CREATE TABLE salmon_farms (id INT, name TEXT, country TEXT); CREATE TABLE temperature_readings (id INT, farm_id INT, temperature FLOAT); INSERT INTO salmon_farms (id, name, country) VALUES (1, 'Farm X', 'Norway'), (2, 'Farm Y', 'Norway'), (3, 'Farm Z', 'Canada'); INSERT INTO temperature_readings (id, farm_id, temperature) VALUES (1, 1, 12.5), (2, 1, 13.0), (3, 2, 11.0), (4, 2, 11.5), (5, 3, 7.0);\n",
      "Sql: SELECT AVG(temperature) FROM temperature_readings TR JOIN salmon_farms SF ON TR.farm_id = SF.id WHERE SF.country = 'Norway';\n",
      "\n",
      "---\n",
      "\n",
      "Follow the following format.\n",
      "\n",
      "Sql Prompt: Natural language query\n",
      "\n",
      "Sql Context: Context for the query\n",
      "\n",
      "Reasoning: Let's think step by step in order to ${produce the sql}. We ...\n",
      "\n",
      "Sql: SQL query\n",
      "\n",
      "---\n",
      "\n",
      "Sql Prompt: Insert a new contractor with \"contractor_id\" 1001, \"name\" \"ABC Construction\", \"location\" \"New York, NY\", and \"license_number\" \"1234567890\" into the \"Contractors\" table.\n",
      "\n",
      "Sql Context: CREATE TABLE Contractors (contractor_id INT, name VARCHAR(255), location VARCHAR(255), license_number VARCHAR(50));\n",
      "\n",
      "Reasoning: Let's think step by step in order to ${insert the new contractor into the Contractors table}. We will first identify the necessary columns and their corresponding values for the new contractor, and then use the INSERT INTO statement to add the new record to the table.\n",
      "\n",
      "Sql: INSERT INTO Contractors (contractor_id, name, location, license_number) VALUES (1001, 'ABC Construction', 'New York, NY', '1234567890');\n",
      "\n",
      "---\n",
      "\n",
      "Sql Prompt: Show the number of vegan hair care products launched in Canada each year.\n",
      "\n",
      "Sql Context: CREATE TABLE hair_care_launches (id INT, product VARCHAR(50), launch_date DATE, country VARCHAR(50), vegan BOOLEAN); INSERT INTO hair_care_launches (id, product, launch_date, country, vegan) VALUES (1, 'Vegan Shampoo', '2021-02-15', 'Canada', TRUE);\n",
      "\n",
      "Reasoning: Let's think step by step in order to produce the SQL query. We need to find the number of vegan hair care products launched in Canada each year.\n",
      "\n",
      "Sql: SELECT YEAR(launch_date) AS launch_year, COUNT(*) AS vegan_products_launched FROM hair_care_launches WHERE country = 'Canada' AND vegan = TRUE GROUP BY launch_year ORDER BY launch_year;\n",
      "\n",
      "---\n",
      "\n",
      "Sql Prompt: List the top 5 countries with the highest number of satellites in orbit as of 2022-01-01, ordered by the number of satellites in descending order.\n",
      "\n",
      "Sql Context: CREATE TABLE countries(id INT, name VARCHAR(255), population INT, satellites_in_orbit INT, last_census_date DATE);\n",
      "\n",
      "Reasoning: Let's think step by step in order to produce the SQL query. We need to find the top 5 countries with the highest number of satellites in orbit as of 2022-01-01, ordered by the number of satellites in descending order.\n",
      "\n",
      "Sql:\u001b[32m SELECT name, satellites_in_orbit\n",
      "FROM countries\n",
      "WHERE last_census_date <= '2022-01-01'\n",
      "ORDER BY satellites_in_orbit DESC\n",
      "LIMIT 5;\u001b[0m\n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "lm.inspect_history(n=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Evaluate the optimized program"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #E0F8E0; padding: 10px; border-radius: 5px;\"> <p style=\"color: #006400; font-size: 18px; font-weight: bold; margin: 0;\"> ✅ Evaluation Stage 3 </p> <p style=\"color: #006400; font-size: 16px; margin: 5px 0 0;\"> With <strong>BootstrapFewShotWithRandomSearch</strong> optimization, <strong>Starling7B</strong> achieves an <strong>80% correctness in validation (25 samples)</strong> and <strong>68% correctness in test (75 samples).</strong> </p> </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starling7b + BootstrapFewShotWithRandomSearch - Validation Score: \n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": []
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Average Metric: 20 / 25  (80.0): 100%|██████████| 25/25 [00:20<00:00,  1.24it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Metric: 20 / 25  (80.0%)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "/home/jjmov99/dspy-testing/.venv/lib/python3.11/site-packages/dspy/evaluate/evaluate.py:187: FutureWarning: DataFrame.applymap has been deprecated. Use DataFrame.map instead.\n",
      "  df = df.applymap(truncate_cell)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "80.0"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(\"Starling7b + BootstrapFewShotWithRandomSearch - Validation Score: \\n\")\n",
    "evaluate = Evaluate(devset=valset, metric=correctness_metric, num_threads=10, display_progress=True, display_table=0)\n",
    "evaluate(optmized_program_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starling7b + BootstrapFewShotWithRandomSearch - Test Score: \n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Average Metric: 51 / 75  (68.0): 100%|██████████| 75/75 [01:01<00:00,  1.22it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Metric: 51 / 75  (68.0%)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "/home/jjmov99/dspy-testing/.venv/lib/python3.11/site-packages/dspy/evaluate/evaluate.py:187: FutureWarning: DataFrame.applymap has been deprecated. Use DataFrame.map instead.\n",
      "  df = df.applymap(truncate_cell)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "68.0"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(\"Starling7b + BootstrapFewShotWithRandomSearch - Test Score: \\n\")\n",
    "evaluate = Evaluate(devset=testset, metric=correctness_metric, num_threads=10, display_progress=True, display_table=0)\n",
    "evaluate(optmized_program_2)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}