one_shot_solver.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/wandb/aihackercup/blob/main/one_shot_solver.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"http://wandb.me/logo-im-png\" width=\"400\" alt=\"Weights & Biases\" />\n",
    "\n",
    "# W&B Lighting Competition - AI Hacker Cup \n",
    "\n",
    "</a>\n",
    "\n",
    "[Weights & Biases](https://wandb.ai/site?utm_source=colab&utm_medium=code&utm_campaign=lightning-ai-hacker-cup) are running a 7-day Lightning Competition focussed on solving practice problems for the  [2024 NeurIPS AI Hacker Cup](https://hackercupai.github.io/) challenge.\n",
    "\n",
    "#### Goal\n",
    "The goal is to try and solve all 5 of the 2023 practice questions for the AI Hacker Cup using MistralAI's models. We’re offering free MistralAI api access via the code in this colab to get people started.\n",
    "\n",
    "#### Competition GitHub\n",
    "The competition [repo here](https://github.com/wandb/aihackercup) contains this colab, the code for the Code Generation Agent and the details on how to make a submission and the competition rules. Note that to run this notebook you'll need to be running it with a T4 GPU (15GB) or larger as the embedding model is run locally.\n",
    "\n",
    "#### Discord\n",
    "You can join the official NeurIPS AI Hacker Cup [discord here](discord.gg/wWeN9hTH32) to share ideas and discuss winning solutions.\n",
    "\n",
    "## Prizes\n",
    "\n",
    "Weights & Biases are giving away a pair of Meta Ray-Ban Smart Glasses for the first individual to submit code that solves:\n",
    "- 3 out of 5 correct solutions\n",
    "- 4 out of 5 correct solutions\n",
    "- 5 out of 5 correct solutions\n",
    "\n",
    "(i.e. in total 3 pairs of sunglasses to give away)\n",
    "\n",
    "## Entry Submissions, Rules & Deadline\n",
    "\n",
    "See the [competition README](https://github.com/wandb/aihackercup) for how to make a submissions the the competition rules."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## W&B Weave\n",
    "\n",
    "[W&B Weave](https://weave-docs.wandb.ai/tutorial-eval?utm_source=colab&utm_medium=code&utm_campaign=lightning-ai-hacker-cup) is used in this competition to run the evaluations. It is a lightweight toolkit for tracking and evaluating LLM applications, built by Weights & Biases. \n",
    "\n",
    "<img src=\"https://raw.githubusercontent.com/wandb/weave/master/docs/static/img/evals-hero.png\" width=\"800\" height=\"450\">\n",
    "\n",
    "If you want to learn more about Weave, you can [get started](https://weave-docs.wandb.ai/quickstart?utm_source=colab&utm_medium=code&utm_campaign=lightning-ai-hacker-cup) by decorating Python functions with `@weave.op`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A simple one-shot solver for the AI Hacker Cup 2024 Qualification Round"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note: You need to run this cell only once**\n",
    "We will clone the starter-kits repo\n",
    "Set the rag folder as our working directory\n",
    "and install the dependencies for the project.\n",
    "\n",
    "**You can comment out the cell after you have run it once.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clone the starter-kits repo\n",
    "!git clone https://github.com/wandb/aihackercup\n",
    "# Change directory to the rag folder. Running the next line twice in the same session will raise an error.\n",
    "%cd aihackercup\n",
    "# Install dependencies\n",
    "!pip install -r requirements.txt -qq"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To run this colab, create a [free Weights & Biases (W&B) account here](https://wandb.ai/site?utm_source=colab&utm_medium=code&utm_campaign=lightning-ai-hacker-cup) and then copy your API key from https://wandb.ai/authorize into the input box below when requested."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import weave\n",
    "\n",
    "WEAVE_PROJECT = \"ai-hacker-cup\"\n",
    "weave_client = weave.init(WEAVE_PROJECT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Select MistralAI models used depending if you want a fast or strong LLM\n",
    "# You can see the full range of MistralAI models here: https://docs.mistral.ai/getting-started/models/\n",
    "FAST_LLM = \"open-mistral-nemo-2407\"\n",
    "STRONG_LLM = \"mistral-large-latest\"\n",
    "\n",
    "os.environ[\"FAST_LLM\"] = STRONG_LLM  # We'll use stong model everywhere\n",
    "os.environ[\"STRONG_LLM\"] = STRONG_LLM\n",
    "\n",
    "# URL for the MistralAI api we'll be using\n",
    "os.environ[\"BASE_URL\"] = \"http://195.242.25.198:8000/v1\"\n",
    "os.environ[\"API_KEY\"] = \"dummy_key\"\n",
    "\n",
    "# Set the max tokens for the models and how many parallel requests to make in Weave Evaluations\n",
    "os.environ[\"MAX_TOKENS\"] = \"4096\"\n",
    "os.environ[\"WEAVE_PARALLELISM\"] = \"2\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "import logging\n",
    "\n",
    "# Start of workout\n",
    "from utils import Problem, async_client, format_response, check_correctness"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Challenges Dataset\n",
    "We will use the **practice** dataset from the **2023** [HackerCup dataset](https://huggingface.co/datasets/hackercupai/hackercup).\n",
    "\n",
    "We have already processed the dataset and saved it as a [`weave.Dataset`](https://weave-docs.wandb.ai/guides/core-types/datasets/?utm_source=colab&utm_medium=code&utm_campaign=lightning-ai-hacker-cup). You can either use the Dataset by running the next cell or download the dataset using the instructions below.\n",
    "\n",
    "We will use this challenge dataset to load some practice problems and solutions from the HackerCup dataset and evaluate our agents on it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get dataset\n",
    "practice_dataset_uri = \"weave:///parambharat/hackercup/object/practice_dataset:R35fXf9N3FE2IOesg7bRPaPAxiE9YbpirhXO9HcHs8w\"\n",
    "problems_dataset = weave.ref(practice_dataset_uri).get().rows[:]\n",
    "problems = list(map(lambda x: Problem(**x), problems_dataset))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's define what we expect as a solution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel, Field\n",
    "\n",
    "class Solution(BaseModel):\n",
    "    core_question: str = Field(..., description=\"Core question of the problem\")\n",
    "    problem_solving_info: str = Field(..., description=\"Problem-solving information related to the core question\")\n",
    "    plan: str = Field(..., description=\"Step by step plan to solve the problem\")\n",
    "    pseudocode: str = Field(..., description=\"Pseudocode to solve the problem\")\n",
    "    source_code: str = Field(..., description=\"Valid Python3 sourcecode to solve the problem.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## One Shot Solver\n",
    "\n",
    "Here we define the One Shot Solver pipeline which:\n",
    "- takes a problem as input\n",
    "- generates a solution using a large language model\n",
    "- executes the generated code\n",
    "- checks if the executed code produces the correct output\n",
    "- returns the solution and test report\n",
    "The solver uses a system prompt and template to guide the LLM in generating\n",
    "a step-by-step solution, including core question extraction, problem-solving plan,\n",
    "pseudocode, and final Python code.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_prompt = \"\"\"\n",
    "You are a world-class competitive programmer tasked with solving a programming problem. \n",
    "You will be provided with a problem statement, and you need to create a Python3 solution for it. \n",
    "Your task it to develop a winning solution to the problem in Python3 programming language.\n",
    "You will do this in a step-by-step manner.\n",
    "\n",
    "Step 1: Extract the core question and the problem-solving information from the problem statement.\n",
    "Step 2: Generate a step by step plan to solve the problem.\n",
    "Step 3: Generate the pseudocode to solve the problem.\n",
    "Step 4: Write the final solution in Python3 programming language to solve the problem.\n",
    "\n",
    "Competition Guidelines:\n",
    "    a. Do not use any external libraries; stick to Python 3 standard library\n",
    "    b. Handle input and output using standard input/output (stdin/stdout)\n",
    "    c. Use helper functions to improve readability of the code.\n",
    "    c. Use the `input()` function to take input from stdin and print the output to stdout.\n",
    "    d. Do not add extra print statements otherwise it will fail the test cases.\n",
    "    e. Make sure your code passes all potential test cases, including edge cases\n",
    "    f. Follow the input/output format specified in the problem statement and the sample test cases.\"\"\"\n",
    "\n",
    "prompt_template = \"\"\"\n",
    "Let's think step by step to solve the problem:\n",
    "\n",
    "Problem: \n",
    "{problem_description}\n",
    "\n",
    "Input: \n",
    "{sample_input}\n",
    "\n",
    "Output: \n",
    "{sample_output}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "@weave.op\n",
    "async def one_shot_solver(\n",
    "    problem: Problem, \n",
    "    llm_model: str,\n",
    "    system_prompt: str, \n",
    "    prompt_template: str,\n",
    "    temperature: float = 0.7,\n",
    "    timeout: int = 10\n",
    ") -> str:\n",
    "    logging.info(f\"Solving problem: {problem.problem_name}\")\n",
    "\n",
    "    # call model one first time to get the code\n",
    "    logging.info(\"Calling model to solve the problem\")\n",
    "    model_output = await async_client.chat.completions.create(\n",
    "        model=llm_model,\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": system_prompt},\n",
    "            {\"role\": \"user\", \"content\": prompt_template.format(\n",
    "                problem_description=problem.problem_description,\n",
    "                sample_input=problem.sample_input,\n",
    "                sample_output=problem.sample_output)}\n",
    "                ],\n",
    "        temperature=temperature,\n",
    "        response_model=None\n",
    "    )\n",
    "\n",
    "    out = model_output.choices[0].message.content\n",
    "\n",
    "    # extract code from the response\n",
    "    logging.info(\"Formatting the response\")\n",
    "    solution = await format_response(out, Solution)\n",
    "\n",
    "    # check if the code is correct\n",
    "    logging.info(\"Checking if the code is correct\")\n",
    "    test_report = await check_correctness(\n",
    "        solution.source_code,\n",
    "        problem.sample_input,\n",
    "        problem.sample_output,\n",
    "        timeout=timeout,\n",
    "    )\n",
    "\n",
    "    return {\"solution\": solution, \"test_report\": test_report}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we are ready to evaluate against the expected solutions.\n",
    "\n",
    "### Create a Weave Model\n",
    "First we create a Weave [\"Model\"](https://weave-docs.wandb.ai/guides/core-types/models?utm_source=colab&utm_medium=code&utm_campaign=lightning-ai-hacker-cup), which has a `predict` function that Weave Evaluations will call to generate a solution. It also has various attributes that we can set to adjust the behaviour of our pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class OneShotSolver(weave.Model):\n",
    "    code_execution_timeout: int = 30\n",
    "    llm_model: str = STRONG_LLM\n",
    "    system_prompt: str = system_prompt\n",
    "    prompt_template: str = prompt_template\n",
    "    temperature: float = 0.7\n",
    "\n",
    "    @weave.op\n",
    "    async def predict(self, problem: dict):\n",
    "        return await one_shot_solver(\n",
    "            problem=Problem(**problem), \n",
    "            llm_model=self.llm_model,\n",
    "            system_prompt=self.system_prompt, \n",
    "            prompt_template=self.prompt_template, \n",
    "            timeout=self.code_execution_timeout,\n",
    "            temperature=self.temperature\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create the Evals Dataset and a Scorer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We expect the output of the \"test_report\" from our agent above to be `\"passed\"` if the solution is correct. You can think of `expected_result` in the `evals_dataset` as the label that the `test_report` from our solver needs to return in order to ensure the generated solution is correct. In this case the scoring is actually happening in our agentic pipeline as the agent needs to know the result so it can decide whether or not to retry.\n",
    "\n",
    "Weave Evaluations expects data formatted as a list of dictionaries for the evaluation dataset. We dump `problem` as a dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evals_dataset = [{\"problem\": problem.model_dump(), \"expected_result\": \"passed\"} for problem in problems]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Weave Evaluations use a scorer function that returns a metric and its result in a dict. Here we define a metric that checks if the code generated by agent passed the test case"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@weave.op\n",
    "def scorer(expected_result: str, model_output: dict) -> dict:\n",
    "    if model_output is None or model_output[\"test_report\"].status is None:\n",
    "        return {\"solution_passed\": False}\n",
    "    return {\"solution_passed\": expected_result == model_output[\"test_report\"].status} # check if the test_report status == passed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = OneShotSolver()\n",
    "\n",
    "evaluator = weave.Evaluation(dataset=evals_dataset, scorers=[scorer], trials=1)\n",
    "\n",
    "results = await evaluator.evaluate(model)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "include_colab_link": true,
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}