colabs/Step1_SolveQueries.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/automix-llm/automix/blob/main/colabs/Step1_SolveQueries.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kIsLu6t_AeDB"
      },
      "source": [
        "# AutoMix: Solving the task\n",
        "\n",
        "- This is the first step of the process. We run inference on both the 13b and 70b models for all tasks. Note that in practice, we don't have to run inference on both the models. This is just for ease of implementation.\n",
        "\n",
        "- Step 2 is verification. Please see the notebook [here](llama13b_f1).\n",
        "\n",
        "\n",
        "*Note: The outputs of this step are provided [here](https://drive.google.com/file/d/1dhyt7UuYumk9Gae9eJ_mpTVrLeSTuRht/view?usp=sharing).*"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LCohyStx1JqL"
      },
      "outputs": [],
      "source": [
        "import pandas as pd\n",
        "from concurrent.futures import ThreadPoolExecutor\n",
        "from functools import partial\n",
        "from tqdm import tqdm\n",
        "from transformers import AutoTokenizer\n",
        "\n",
        "# Using the tokenizer\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"meta-llama/Llama-2-13b-hf\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yQ0TiRTm25lp"
      },
      "source": [
        "### Read data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XBoOaj7bcSCW"
      },
      "outputs": [],
      "source": [
        "# get the input file from https://drive.google.com/file/d/1dhyt7UuYumk9Gae9eJ_mpTVrLeSTuRht/view?usp=sharing"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9TrDKNY-1Nt9"
      },
      "outputs": [],
      "source": [
        "inputs = pd.read_json(\"data/automix_input.jsonl\", lines=True, orient=\"records\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Oez36MjOBxoz",
        "outputId": "ac39e751-bb66-4ce0-d09e-13de2a059eba"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Number of inputs: 10\n"
          ]
        }
      ],
      "source": [
        "demo = True #@param {type:\"boolean\"}\n",
        "if demo:\n",
        "    inputs = inputs.sample(10)\n",
        "print(f\"Number of inputs: {len(inputs)}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 81
        },
        "id": "LhZITST3c8Us",
        "outputId": "789aee9c-ef3b-44c1-b53b-bdcc6b216ee9"
      },
      "outputs": [
        {
          "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>id</th>\n",
              "      <th>pid</th>\n",
              "      <th>base_ctx</th>\n",
              "      <th>question</th>\n",
              "      <th>output</th>\n",
              "      <th>dataset</th>\n",
              "      <th>split</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>24193</th>\n",
              "      <td>bb11261719535ef6b4f9e092be690c99861e7f0a4aeabc...</td>\n",
              "      <td>58402da9f8baad0b4bb42810289b8a64e7f707e40b4149...</td>\n",
              "      <td>(CNN) -- Three Pakistani cricketers found guil...</td>\n",
              "      <td>how did they plead?</td>\n",
              "      <td>guilty</td>\n",
              "      <td>coqa</td>\n",
              "      <td>train</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "                                                      id  \\\n",
              "24193  bb11261719535ef6b4f9e092be690c99861e7f0a4aeabc...   \n",
              "\n",
              "                                                     pid  \\\n",
              "24193  58402da9f8baad0b4bb42810289b8a64e7f707e40b4149...   \n",
              "\n",
              "                                                base_ctx             question  \\\n",
              "24193  (CNN) -- Three Pakistani cricketers found guil...  how did they plead?   \n",
              "\n",
              "       output dataset  split  \n",
              "24193  guilty    coqa  train  "
            ]
          },
          "execution_count": 49,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "inputs.head(1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1Chiv4DP28Yz",
        "outputId": "913957a9-b80a-495e-836a-8d1bd58b17e7"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "cnli            3\n",
              "quality         3\n",
              "coqa            2\n",
              "narrative_qa    1\n",
              "qasper          1\n",
              "Name: dataset, dtype: int64"
            ]
          },
          "execution_count": 50,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "inputs['dataset'].value_counts()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_BacMsdv29AT"
      },
      "source": [
        "### Run inference on each task\n",
        "- The prompts are taken from [zero-scrolls](https://www.zero.scrolls-benchmark.com/)\n",
        "- For dataset construction, please see the paper. TLDR: narrative qa, qasper, quality, and cnli are taken from [scrolls](https://www.scrolls-benchmark.com/), and coqa is from [huggingface](https://huggingface.co/datasets/coqa).\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W-VdTgsXEvnd"
      },
      "source": [
        "#### Task prompts"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hRLXnjM02Z3d"
      },
      "outputs": [],
      "source": [
        "dataset_prompts_and_instructions = {\n",
        "\n",
        "### NARRATIVE_QA\n",
        "\n",
        "    \"narrative_qa\": {\n",
        "        \"instruction\": \"You are given a story, which can be either a novel or a movie script, and a question. Answer the question as concisely as you can, using a single phrase if possible.\",\n",
        "        \"prompt\": \"\"\"Story:\n",
        "{context}\n",
        "\n",
        "{instruction}\n",
        "\n",
        "Question: {question}\n",
        "\n",
        "Answer: The answer is'\"\"\",\n",
        "        \"truncation_message\": \"... [The rest of the story is omitted]\\n\\n\",\n",
        "    },\n",
        "\n",
        "### QASPER\n",
        "\n",
        "    \"qasper\": {\n",
        "        \"instruction\": \"You are given a scientific article and a question. Answer the question as concisely as you can, using a single phrase or sentence if possible. If the question cannot be answered based on the information in the article, write 'unanswerable'. If the question is a yes/no question, answer 'yes', 'no', or 'unanswerable'.\",\n",
        "        \"prompt\": \"\"\"Article:\n",
        "{context}\n",
        "\n",
        "{instruction}\n",
        "\n",
        "Question: {question}\n",
        "\n",
        "Answer: The answer is'\"\"\",\n",
        "        \"truncation_message\": \"... [The rest of the article is omitted]\\n\\n\",\n",
        "    },\n",
        "\n",
        "### QUALITY\n",
        "\n",
        "\n",
        "\"quality\": {\n",
        "        \"instruction\": \"You are provided a story and a multiple-choice question with 4 possible answers (marked by A, B, C, D). Choose the best answer by writing its corresponding letter (either A, B, C, or D).\",\n",
        "        \"prompt\": \"\"\"Story:\n",
        "{context}\n",
        "\n",
        "{instruction}\n",
        "\n",
        "Question and Possible Answers: {question}\n",
        "\n",
        "Answer: The answer is'\"\"\",\n",
        "\n",
        "        \"truncation_message\": \"... [The rest of the story is omitted]\\n\\n\",\n",
        "    },\n",
        "\n",
        "\n",
        "### CNLI\n",
        "\n",
        "    \"cnli\": {\n",
        "        \"instruction\": \"You are given a non-disclosure agreement and a sentence that proposes a hypothesis based on the agreement. Choose whether the hypothesis is entailed by the agreement, contradicted by the agreement, or not mentioned by (neutral to) the agreement. If the hypothesis is entailed by the agreement, write 'Entailment'. If the hypothesis is contradicted by the agreement, write 'Contradiction'. If the hypothesis is not mentioned by the agreement, write 'Not mentioned'.\",\n",
        "        \"prompt\": \"\"\"Contract:\n",
        "{context}\n",
        "\n",
        "{instruction}\n",
        "\n",
        "Hypothesis: {question}\n",
        "\n",
        "Answer: The answer is'\"\"\",\n",
        "        \"truncation_message\": \"... [The rest of the contract is omitted]\\n\\n\",\n",
        "    },\n",
        "\n",
        "### COQA\n",
        "\n",
        "      \"coqa\": {\n",
        "        \"instruction\": \"You are given a story, which can be either a novel or a movie script, and a question. Answer the question as concisely as you can, using a single phrase if possible.\",\n",
        "        \"prompt\": \"\"\"Story:\n",
        "{context}\n",
        "\n",
        "{instruction}\n",
        "\n",
        "Question: {question}\n",
        "\n",
        "Answer: The answer is'\"\"\",\n",
        "        \"truncation_message\": \"... [The rest of the story is omitted]\\n\\n\",\n",
        "    },\n",
        "\n",
        "\n",
        "}\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "z8tcXajF2_Aq"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2efylito2_cK"
      },
      "source": [
        "#### LLM tooling setup\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lpCw0K9e3Bto"
      },
      "outputs": [],
      "source": [
        "import openai\n",
        "\n",
        "openai.api_key = \"EMPTY\"\n",
        "openai.api_base = \"http://pitt.lti.cs.cmu.edu:8003/v1\"\n",
        "\n",
        "def call_openai_api(\n",
        "    prompt: str,\n",
        "    engine_name: str,\n",
        "    temperature: float = 0.0,\n",
        "    n: int = 1,\n",
        "    stop: str = '\\n',\n",
        "    max_tokens: int = 100,\n",
        "    batch_size: int = 32\n",
        "):\n",
        "    \"\"\"\n",
        "    Call the OpenAI API to create completions based on the input prompt.\n",
        "\n",
        "    Parameters:\n",
        "    - prompt (str): The prompt.\n",
        "    - engine_name (str, optional): The engine to use for the completion.\n",
        "    - temperature (float, optional): Sampling temperature for randomness. Defaults to 0.0.\n",
        "    - n (int, optional): Number of completions to generate. Defaults to 1.\n",
        "    - stop (str, optional): Token at which the API should stop generating further tokens. Defaults to '\\n'.\n",
        "    - max_tokens (int, optional): Maximum number of tokens in the generated output. Defaults to 100.\n",
        "    - batch_size (int, optional): Maximum num_completions for each API call. Defaults to 32.\n",
        "\n",
        "    Returns:\n",
        "    - list/str: Generated text completions from the API. Returns a list of strings if n > 1, else a single string.\n",
        "    \"\"\"\n",
        "    all_responses = []\n",
        "    orig_n = n\n",
        "\n",
        "    try:\n",
        "        while n > 0:\n",
        "            current_batch_size = min(n, batch_size)\n",
        "\n",
        "            response = openai.Completion.create(\n",
        "                        model=engine_name,\n",
        "                        prompt=prompt,\n",
        "                        temperature=temperature,\n",
        "                        max_tokens=max_tokens,\n",
        "                        n=current_batch_size,\n",
        "                        stop=stop,\n",
        "                    )\n",
        "\n",
        "            all_responses.extend([choice['text'] for choice in response['choices']])\n",
        "\n",
        "            n -= current_batch_size\n",
        "\n",
        "        return all_responses if orig_n > 1 else all_responses[0]\n",
        "\n",
        "    except Exception as e:\n",
        "        print(f\"An error occurred: {str(e)}\")\n",
        "        return None\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b1UkBG6g2y2T"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NmE-nMV7dzh4"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0LP86l4pEyns"
      },
      "source": [
        "#### Run inference"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8Bg2l6hh1t8b",
        "outputId": "477fc90e-9ec3-46c3-c114-02a02ac7277f"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████| 10/10 [00:09<00:00,  1.08it/s]\n",
            "100%|██████████| 10/10 [00:09<00:00,  1.03it/s]\n"
          ]
        }
      ],
      "source": [
        "\n",
        "def run_solver_job(df, prepare_row_func, engine_name: str, max_workers: int = 32,\n",
        "                   temperature: float = 0.0, n: int = 1, stop: str = '\\n',\n",
        "                   max_tokens: int = 100):\n",
        "    \"\"\"\n",
        "    Runs a solver job using a specified engine, applying concurrent futures and tqdm for progress tracking.\n",
        "\n",
        "    Parameters:\n",
        "    - df: Input dataframe\n",
        "    - prepare_row_func: Function to prepare rows of df for the solver\n",
        "    - engine_name (str): Name of the engine to use\n",
        "    - max_workers (int, optional): Maximum number of workers for ThreadPoolExecutor. Defaults to 32.\n",
        "    - temperature (float, optional): Temperature parameter for call_openai_api. Defaults to 0.0.\n",
        "    - n (int, optional): n parameter for call_openai_api. Defaults to 1.\n",
        "    - stop (str, optional): Stop parameter for call_openai_api. Defaults to '\\n'.\n",
        "    - max_tokens (int, optional): Maximum number of tokens for call_openai_api. Defaults to 100.\n",
        "\n",
        "    Returns:\n",
        "    - list: Results from the solver job\n",
        "    \"\"\"\n",
        "    # Creating a partial function with specified parameters\n",
        "    solver_call = partial(call_openai_api,\n",
        "                          engine_name=engine_name,\n",
        "                          temperature=temperature,\n",
        "                          n=n,\n",
        "                          stop=stop,\n",
        "                          max_tokens=max_tokens)\n",
        "\n",
        "    # Running the solver job concurrently and tracking progress with tqdm\n",
        "    with ThreadPoolExecutor(max_workers=max_workers) as executor:\n",
        "        results = list(tqdm(executor.map(solver_call, df.apply(prepare_row_func, axis=1)),\n",
        "                            total=df.shape[0]))\n",
        "\n",
        "    return results\n",
        "\n",
        "def prepare_row(row):\n",
        "    dataset = row[\"dataset\"]\n",
        "    prompt = dataset_prompts_and_instructions[dataset][\"prompt\"]\n",
        "    instruction = dataset_prompts_and_instructions[dataset][\"instruction\"]\n",
        "    question = row['question']\n",
        "    context = row['base_ctx']\n",
        "\n",
        "    full_text = prompt.format(context=context, instruction=instruction, question=question)\n",
        "\n",
        "    tokens = tokenizer.encode(full_text)\n",
        "\n",
        "    # Check if the length exceeds 3096 tokens, llama2 requirements\n",
        "    if len(tokens) > 3096:\n",
        "        tokens = tokens[-3096:]\n",
        "\n",
        "    truncated_text = tokenizer.decode(tokens)\n",
        "\n",
        "    return truncated_text\n",
        "\n",
        "\n",
        "# Engine names for \"13b\" and \"70b\" - Replace these with the actual engine names.\n",
        "engine_13b = \"meta-llama/Llama-2-13b-hf\"\n",
        "engine_70b = \"meta-llama/Llama-2-70b-hf\"\n",
        "\n",
        "\n",
        "# Running the job for both engines and storing results.\n",
        "# Note that we use temp = 0.0 for task.\n",
        "\n",
        "results_13b = run_solver_job(inputs, prepare_row, engine_13b)\n",
        "results_70b = run_solver_job(inputs, prepare_row, engine_70b)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xvpggnDyAFDb",
        "outputId": "6d57648c-dc04-4a72-f19c-11a486afa933"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "\"slop bucket'.\""
            ]
          },
          "execution_count": 53,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "results_70b[0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gScf9qD-kCc1"
      },
      "outputs": [],
      "source": [
        "def clean_answer(ans: str) -> str:\n",
        "  return ans.replace(\"'\", \"\") if ans else pd.NA"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MV8Z5yfi3SAx",
        "outputId": "007e9853-7f42-4a92-bc2a-598251e32f66"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "10"
            ]
          },
          "execution_count": 55,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "len(results_70b)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6xPfynOSwO4X"
      },
      "outputs": [],
      "source": [
        "inputs['llama13b_pred_ans'] = [clean_answer(ans) for ans in results_13b]\n",
        "inputs['llama70b_pred_ans'] = [clean_answer(ans) for ans in results_70b]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ai6AHTgTwWw4"
      },
      "outputs": [],
      "source": [
        "inputs_with_predictions =  inputs.dropna()\n",
        "# slighly better name for inputs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ctpJ4oFTjZrF",
        "outputId": "59ff3096-58f6-48f8-c8b5-966b3fa8ffc3"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "10/10 inputs have predictions\n"
          ]
        }
      ],
      "source": [
        "print(f\"{len(inputs_with_predictions)}/{len(inputs)} inputs have predictions\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "l0vm6-Y9wk3X"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pnrEPWt3gzUo"
      },
      "source": [
        "## Add scores"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pL7siF5ni8tc"
      },
      "outputs": [],
      "source": [
        "import re\n",
        "import string\n",
        "from collections import Counter\n",
        "import pandas as pd\n",
        "\n",
        "def f1_score(prediction, ground_truth):\n",
        "    prediction_tokens = normalize_answer(prediction).split()\n",
        "    ground_truth_tokens = normalize_answer(ground_truth).split()\n",
        "    common = Counter(prediction_tokens) & Counter(ground_truth_tokens)\n",
        "    num_same = sum(common.values())\n",
        "    if num_same == 0:\n",
        "        return 0\n",
        "    precision = 1.0 * num_same / len(prediction_tokens)\n",
        "    recall = 1.0 * num_same / len(ground_truth_tokens)\n",
        "    f1 = (2 * precision * recall) / (precision + recall)\n",
        "    return f1\n",
        "\n",
        "def normalize_answer(s):\n",
        "    \"\"\"Lower text and remove punctuation, articles, and extra whitespace.\"\"\"\n",
        "    def remove_articles(text):\n",
        "        return re.sub(r\"\\b(a|an|the)\\b\", \" \", text)\n",
        "\n",
        "    def white_space_fix(text):\n",
        "        return \" \".join(text.split())\n",
        "\n",
        "    def remove_punc(text):\n",
        "        exclude = set(string.punctuation)\n",
        "        return \"\".join(ch for ch in text if ch not in exclude)\n",
        "\n",
        "    def lower(text):\n",
        "        return text.lower()\n",
        "\n",
        "    return white_space_fix(remove_articles(remove_punc(lower(s))))\n",
        "\n",
        "def calculate_f1_for_models(df, model_sizes, ground_truth_col='output'):\n",
        "    \"\"\"\n",
        "    Calculates F1 score for different model sizes and adds the results as new columns in the DataFrame.\n",
        "\n",
        "    Parameters:\n",
        "    - df (pd.DataFrame): The DataFrame containing prediction data.\n",
        "    - model_sizes (list of str): List containing strings that denote model sizes.\n",
        "      Used to create column names dynamically.\n",
        "    - ground_truth_col (str, optional): The name of the column containing ground truth data.\n",
        "      Defaults to 'output'.\n",
        "\n",
        "    Returns:\n",
        "    - pd.DataFrame: The original DataFrame with added columns for the F1 scores.\n",
        "    \"\"\"\n",
        "    for size in model_sizes:\n",
        "        pred_col = f'llama{size}_pred_ans'\n",
        "        f1_col = f'llama{size}_f1'\n",
        "        df[f1_col] = df.apply(\n",
        "            lambda r: f1_score(prediction=r[pred_col], ground_truth=r[ground_truth_col]),\n",
        "            axis=1\n",
        "        )\n",
        "    return df\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_ILypFipBPyp"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pfX0D0d2BQBV"
      },
      "source": [
        "#### For quality, LLAMA2-13b sometimes generates only the option (e.g., a). Simple matching with output won't work, so we have to do map the generated option to the correct answer and do the matching.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SQF3lddGkl82"
      },
      "outputs": [],
      "source": [
        "import pandas as pd\n",
        "import re\n",
        "from typing import List\n",
        "\n",
        "def extract_option(row: pd.Series) -> str:\n",
        "    \"\"\"\n",
        "    Extracts the correct option from the provided row.\n",
        "\n",
        "    Parameters:\n",
        "        row (pd.Series): A row of a DataFrame, expected to contain 'question' and 'output' columns.\n",
        "\n",
        "    Returns:\n",
        "        str: The letter of the correct option, or None if not found.\n",
        "    \"\"\"\n",
        "    options = re.findall(r'\\((\\w)\\) ([\\w\\s]+)', row['question'])\n",
        "    for option, value in options:\n",
        "        if value.strip() == row['output'].strip():\n",
        "            return option\n",
        "    return None\n",
        "\n",
        "def extract_option_from_prediction(pred: str) -> str:\n",
        "    \"\"\"\n",
        "    Extracts the selected option letter from a prediction string.\n",
        "\n",
        "    Parameters:\n",
        "        pred (str): The prediction string, expected to start with an option letter.\n",
        "\n",
        "    Returns:\n",
        "        str: The extracted option letter, or None if not found or if `pred` is empty.\n",
        "    \"\"\"\n",
        "    if len(pred.strip()) == 0:\n",
        "        return None\n",
        "\n",
        "    option = pred.split()[0]\n",
        "    for char in option:\n",
        "        if char in ['A', 'B', 'C', 'D']:\n",
        "            return char\n",
        "    return None\n",
        "\n",
        "def calculate_f1_for_multi_choice(df: pd.DataFrame, model_sizes: List[str], datasets: List[str]=[\"quality\"]) -> pd.DataFrame:\n",
        "    \"\"\"\n",
        "    Computes F1 scores for predictions in multiple-choice format.\n",
        "\n",
        "    It extracts correct and predicted options and computes F1 scores, with special handling\n",
        "    for certain datasets. This function mutates the input DataFrame by adding new columns\n",
        "    for extracted options and possibly modifying F1 scores.\n",
        "\n",
        "    Parameters:\n",
        "        df (pd.DataFrame): The DataFrame containing prediction and ground truth data.\n",
        "            Expected to contain columns in the format 'llama{size}_pred_ans'.\n",
        "        model_sizes (List[str]): List of strings indicating the model sizes for which\n",
        "            predictions are available in `df` (e.g., ['13b', '70b']).\n",
        "        datasets (List[str], optional): List of dataset names that require special handling.\n",
        "            Defaults to [\"quality\"].\n",
        "\n",
        "    Returns:\n",
        "        pd.DataFrame: The original DataFrame with additional/modified columns for extracted\n",
        "            options and potentially modified F1 scores.\n",
        "    \"\"\"\n",
        "    df['correct_option'] = df.apply(extract_option, axis=1)\n",
        "\n",
        "    for size in model_sizes:\n",
        "        pred_ans_col = f'llama{size}_pred_ans'\n",
        "        pred_option_col = f'llama{size}_pred_option'\n",
        "        f1_col = f'llama{size}_f1'\n",
        "\n",
        "        # Remove single quotes from predictions for specified datasets\n",
        "        df[pred_ans_col] = df.apply(lambda r: r[pred_ans_col] if r[\"dataset\"] not in datasets else r[pred_ans_col].replace(\"'\", \"\"), axis=1)\n",
        "\n",
        "        # Extract the option from the prediction\n",
        "        df[pred_option_col] = df[pred_ans_col].apply(extract_option_from_prediction)\n",
        "\n",
        "        # Compute the F1 score: if dataset is in `datasets`, F1 is 1 if predicted option matches correct option, else it's 0\n",
        "        df[f1_col] = df.apply(lambda r: r[pred_option_col] == r['correct_option'] if r[\"dataset\"] in datasets else r[f1_col], axis=1)\n",
        "\n",
        "    return df\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "juiEO5iilJJG"
      },
      "outputs": [],
      "source": [
        "model_sizes = ['13b', '70b']\n",
        "\n",
        "# Calculating F1 scores for each model size\n",
        "inputs_with_predictions = calculate_f1_for_models(inputs_with_predictions, model_sizes)\n",
        "\n",
        "# Further processing and calculating F1 scores for multi-choice questions\n",
        "inputs_with_predictions = calculate_f1_for_multi_choice(inputs_with_predictions, model_sizes)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "f5z6bsGMEVeJ",
        "outputId": "3b4aa969-8a54-404e-d2bb-b8664d336822"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "llama13b_f1    0.2\n",
              "llama70b_f1    0.2\n",
              "dtype: float64"
            ]
          },
          "execution_count": 68,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "inputs_with_predictions[['llama13b_f1', 'llama70b_f1']].mean()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Sx9-jpydEXYG"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "authorship_tag": "ABX9TyNP8Py/sLdDPdKd/J5KZ4P5",
      "collapsed_sections": [
        "yQ0TiRTm25lp",
        "W-VdTgsXEvnd"
      ],
      "include_colab_link": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}