{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/cmu-llms-class/11-766-hw1-dev/blob/main/src/task2/11_766_HW1_Problem2_food_classifier_merged.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ic0qyjeoy6Uk"
      },
      "source": [
        "# Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c8qQySFgaYwm"
      },
      "outputs": [],
      "source": [
        "# Install Pytorch & other libraries\n",
        "%pip install \"torch==2.4.1\" tensorboard\n",
        "# flash-attn might need to be deleted due to library version unmatch\n",
        "%pip install flash-attn \"setuptools<71.0.0\" scikit-learn\n",
        "\n",
        "# Install Hugging Face libraries\n",
        "%pip install  --upgrade \\\n",
        "  \"datasets==3.1.0\" \\\n",
        "  \"hf-transfer==0.1.8\" \\\n",
        "  \"transformers==4.48.1\" \\\n",
        "  \"accelerate>=0.34.0\" \\\n",
        "  \"peft==0.13.2\" \\\n",
        "  # \"accelerate==1.2.1\" \\\n",
        "  # \"transformers==4.47.1\"\n",
        "%pip uninstall -y torchvision\n",
        "%pip install --no-cache-dir --index-url https://download.pytorch.org/whl/cu121 torchvision==0.19.1+cu121\n",
        "# ModernBERT is not yet available in an official release, so we need to install it from github\n",
        "# %pip install \"git+https://github.com/huggingface/transformers.git@6e0515e99c39444caae39472ee1b2fd76ece32f1\" --upgrade\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gZqLZXGKQude"
      },
      "outputs": [],
      "source": [
        "# @title Imports\n",
        "import pandas as pd\n",
        "import gspread\n",
        "import pandas as pd\n",
        "from google.colab import userdata\n",
        "from openai import OpenAI\n",
        "from tqdm import tqdm\n",
        "\n",
        "from google.colab import auth\n",
        "auth.authenticate_user()\n",
        "\n",
        "import gspread\n",
        "from google.auth import default\n",
        "creds, _ = default()\n",
        "\n",
        "gc = gspread.authorize(creds)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GdmEk0TMtMfa"
      },
      "source": [
        "# Part 1\n",
        "\n",
        "Make a copy of [this spreadsheet](https://docs.google.com/spreadsheets/d/1rKTpPvXESM5PDOc4clAHiYuowC9zi5P3WbkknWAZJgs/edit?usp=sharing) and annotate the data. Then, update the `SHEET_URL` variable below to link to your copy of the spreadsheet."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N4SoXm9ZxDkV"
      },
      "outputs": [],
      "source": [
        "# @title Load data\n",
        "\n",
        "# SHEET_URL = \"https://docs.google.com/spreadsheets/d/1Aq9NJT7vjbjI-1NJMk125XV6E6Sf-39jxuawtv3A5S4/edit?gid=0#gid=0\" # @param {\"type\":\"string\",\"placeholder\":\"\"}\n",
        "SHEET_URL = \"https://docs.google.com/spreadsheets/d/1Lt_SX9QGHB08XaxermeI7q61TRbwwM6B12HzdpqAf_E\" # @param {\"type\":\"string\",\"placeholder\":\"\"}\n",
        "\n",
        "worksheet = gc.open_by_url(SHEET_URL).sheet1\n",
        "\n",
        "rows = worksheet.get_all_values()\n",
        "df = pd.DataFrame.from_records(rows1:=rows[1:], columns=rows[0])\n",
        "\n",
        "LABEL_NAMES  = [\"wouldn't want to eat\", \"meh\", \"sounds tasty\"]\n",
        "num_labels = len(label_names)\n",
        "\n",
        "train_df = df[df[\"Split\"] == \"train\"]\n",
        "valid_df = df[df[\"Split\"] == \"valid\"]\n",
        "test_df = df[df[\"Split\"] == \"test\"]\n",
        "\n",
        "print(df[\"Label\"].value_counts())\n",
        "valid_df"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cmMdfU5uytrP"
      },
      "source": [
        "# Part 2: Use the OpenAI API to build a food tastiness classifer.\n",
        "\n",
        "To begin, you should create an API key with OpenAI. Then add that key as a secret to Colab by clicking on the key symbol on the left. Give it the name `OPENAI_API_KEY`.\n",
        "\n",
        "We have provided you a very simple classifier implementation that does in-context learning.\n",
        "You should try to improve this classifier using any of the following approaches:\n",
        "\n",
        "- [Finetuning API](https://platform.openai.com/docs/guides/supervised-fine-tuning)\n",
        "- [Structured model outputs](https://platform.openai.com/docs/guides/structured-outputs?api-mode=chat)\n",
        "- Better few-shot prompting\n",
        "- More complex system instruction and prompt format\n",
        "- [Different models](https://platform.openai.com/docs/models) and [inference parameters](https://platform.openai.com/docs/api-reference/responses/create)\n",
        "- Performing one annotation per API call versus several annotations per API call\n",
        "\n",
        "However, please stick with the OpenAI API for this question, and remember to be mindful of how much you are spending. You should try to spend no more than 15 USD during your experimentation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4YzGwChOvRw5"
      },
      "outputs": [],
      "source": [
        "# @title Functions for calling OpenAI API\n",
        "_CLIENT = OpenAI(api_key=userdata.get('OPENAI_API_KEY'))\n",
        "\n",
        "def create_system_prompt(df):\n",
        "  few_shot_examples = [f\"{row.Food}\\t{row.Label}\" for row in df.itertuples()]\n",
        "  few_shot_examples = few_shot_examples\n",
        "  few_shot_examples = \"\\n\".join(few_shot_examples)\n",
        "\n",
        "  preferences = df\n",
        "  s = \"You are an expert predictor of my food preferences. \" +\\\n",
        "      \"For each food item I provide, you should output \\\"1\\\" if you think \" +\\\n",
        "      \"I'll dislike it, \\\"2\\\" if you think I'll have no opinion, and \\\"3\\\"if \" +\\\n",
        "      \"you think I'll really like it.\\n\\n\" +\\\n",
        "      f\"Here are some preferences to start you off: {few_shot_examples}\\n\\n\" +\\\n",
        "      \"Answer with just the number.\"\n",
        "\n",
        "  return s\n",
        "\n",
        "def prompt_gpt5(food_name):\n",
        "  response = _CLIENT.responses.create(\n",
        "    model=\"gpt-5-nano-2025-08-07\",\n",
        "    input=[\n",
        "      {\n",
        "        \"role\": \"developer\",\n",
        "        \"content\": [\n",
        "            {\"type\": \"input_text\", \"text\": create_system_prompt(train_df)}]\n",
        "      },\n",
        "      {\n",
        "        \"role\": \"user\",\n",
        "        \"content\": [{\"type\": \"input_text\", \"text\": food_name}]\n",
        "      },\n",
        "    ],\n",
        "    text={\n",
        "      \"format\": {\n",
        "        \"type\": \"text\"\n",
        "      },\n",
        "      \"verbosity\": \"medium\"\n",
        "    },\n",
        "    reasoning={\n",
        "      \"effort\": \"medium\",\n",
        "      \"summary\": \"auto\"\n",
        "    },\n",
        "    tools=[],\n",
        "    store=True,\n",
        "    include=[\n",
        "      \"reasoning.encrypted_content\",\n",
        "    ]\n",
        "  )\n",
        "  return response\n",
        "\n",
        "\n",
        "system_prompt = create_system_prompt(train_df)\n",
        "print(\"===System prompt===\")\n",
        "print(system_prompt)\n",
        "\n",
        "print(\"\\n===Grilled Cheese Test===\")\n",
        "response = prompt_gpt5(\"Grilled Cheese Sandwich\")\n",
        "print(response.output[-1].content[0].text)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7b_XeY_uUbH7"
      },
      "outputs": [],
      "source": [
        "# @title Run inference on full validation set\n",
        "\n",
        "gpt_responses = []\n",
        "for row in tqdm(valid_df.itertuples(), total=len(valid_df)):\n",
        "  response = prompt_gpt5(row.Food)\n",
        "  gpt_responses.append(response)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "81M_ugV1f9jG"
      },
      "outputs": [],
      "source": [
        "# @title Evaluation\n",
        "predicted_labels = []\n",
        "true_labels = []\n",
        "correct = []\n",
        "\n",
        "for response, row in zip(gpt_responses, valid_df.itertuples()):\n",
        "  predicted_label = int(response.output[-1].content[0].text)\n",
        "  true_label = int(row.Label[0])\n",
        "\n",
        "  predicted_labels.append(predicted_label)\n",
        "  true_labels.append(true_label)\n",
        "  correct.append(predicted_label == true_label)\n",
        "\n",
        "print(\"Accuracy:\", sum(correct) / len(correct))\n",
        "\n",
        "import numpy as np\n",
        "from sklearn.metrics import confusion_matrix\n",
        "import matplotlib.pyplot as plt\n",
        "import seaborn as sns\n",
        "\n",
        "# Generate the confusion matrix\n",
        "cm = confusion_matrix(true_labels, predicted_labels)\n",
        "\n",
        "# Plot the confusion matrix\n",
        "plt.figure(figsize=(8, 6))\n",
        "sns.heatmap(\n",
        "    cm, annot=True, fmt=\"d\", cmap=\"Blues\",\n",
        "    xticklabels=LABEL_NAMES, yticklabels=LABEL_NAMES)\n",
        "plt.xlabel(\"Predicted Label\")\n",
        "plt.ylabel(\"True Label\")\n",
        "plt.title(\"Confusion Matrix\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "R2BeVd7hYQtv"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "from sklearn.metrics import confusion_matrix\n",
        "import matplotlib.pyplot as plt\n",
        "import seaborn as sns\n",
        "\n",
        "# Generate the confusion matrix\n",
        "cm = confusion_matrix(true_labels, predicted_labels)\n",
        "\n",
        "# Plot the confusion matrix\n",
        "plt.figure(figsize=(8, 6))\n",
        "sns.heatmap(cm, annot=True, fmt=\"d\", cmap=\"Blues\",\n",
        "            xticklabels=label_names, yticklabels=label_names)\n",
        "plt.xlabel(\"Predicted Label\")\n",
        "plt.ylabel(\"True Label\")\n",
        "plt.title(\"Confusion Matrix\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xMvc2R-xzeQs"
      },
      "source": [
        "# Part 3: Finetune a tiny model on the synthetically labeled data."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KhFSQ25DaYwg"
      },
      "source": [
        "**Fine-tune classifier with ModernBERT:** Large Language Models (LLMs) have become ubiquitous in 2024. However, smaller, specialized models - particularly for classification tasks - remain critical for building efficient and cost-effective AI systems. In this part, we will fine-tune our smaller models (ModernBERT, a new tiny encoder model) on your own food classification dataset, and evaluate the fine-tuned model on your own food testset.\n",
        "\n",
        "ModernBERT is a refreshed version of BERT models, with 8192 token context length, significantly better downstream performance, and much faster processing speeds.\n",
        "\n",
        "For this part, you will need to:\n",
        "1. Load and prepare the classification dataset  \n",
        "2. Fine-tune & evaluate ModernBERT with the Hugging Face `Trainer`\n",
        "3. Run inference & test model\n",
        "\n",
        "**Quick intro: ModernBERT**: ModernBERT is a modernization of BERT maintaining full backward compatibility while delivering dramatic improvements through architectural innovations like rotary positional embeddings (RoPE), alternating attention patterns, and hardware-optimized design. The model comes in two sizes:\n",
        "- ModernBERT Base (139M parameters, we are using it in this part)\n",
        "- ModernBERT Large (395M parameters)\n",
        "\n",
        "ModernBERT achieves state-of-the-art performance across classification, retrieval and code understanding tasks while being 2-4x faster than previous encoder models. This makes it ideal for high-throughput production applications like LLM routing, where both accuracy and latency are critical.\n",
        "\n",
        "ModernBERT was trained on 2 trillion tokens of diverse data including web documents, code, and scientific articles - making it much more robust than traditional BERT models trained primarily on Wikipedia. This broader knowledge helps it better understand the nuances of user prompts across different domains.\n",
        "\n",
        "If you want to learn more about ModernBERT's architecture and training process, check out the official [blog](https://huggingface.co/blog/modernbert).\n",
        "\n",
        "---\n",
        "\n",
        "Now let's get started building our LLM router with ModernBERT! 🚀\n",
        "\n",
        "*Note: This part was created and tested on a Colab T4 GPU.*\n",
        "\n",
        "### Setup environment and install libraries\n",
        "\n",
        "Our first step is to install Hugging Face Libraries and Pyroch, including transformers and datasets."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8S4Qv__caYwo"
      },
      "source": [
        "We will use the [Hugging Face Hub](https://huggingface.co/models) as a remote model versioning service. This means we will automatically push our model, logs and information to the Hub during training. You must register on the [Hugging Face](https://huggingface.co/join) for this, then create an access token at [this link](https://huggingface.co/settings/tokens).\n",
        "\n",
        "Add your token to the Colab's list of secrets with the name `HF_TOKEN`.\n",
        "\n",
        "We will use the `login` util from the `huggingface_hub` package to log into our account and store our token (access key) on the disk."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jwX0kpVtaYwp"
      },
      "outputs": [],
      "source": [
        "from huggingface_hub import login\n",
        "\n",
        "login(token=userdata.get('HF_TOKEN'), add_to_git_credential=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jJZSVRgGaYwp"
      },
      "source": [
        "### 1. Load and prepare the dataset\n",
        "\n",
        "We will fine-tune ModernBERT on the same food-preference data structure produced in the previous part 3. Concretely:\n",
        "Let's use the [🤗 Datasets](https://huggingface.co/docs/datasets/index) library to build a `DatasetDict` from each DataFrame."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Xqp38ztUpb-d"
      },
      "outputs": [],
      "source": [
        "from datasets import Dataset, DatasetDict\n",
        "\n",
        "def _ensure_int_labels(df):\n",
        "    \"\"\"Extract numeric labels from string format like '1 - wouldn't want to eat'.\"\"\"\n",
        "    df = df[[\"Food\", \"Label\"]].dropna().reset_index(drop=True).copy()\n",
        "\n",
        "    # Handle both string format (\"1 - ...\") and integer format\n",
        "    def extract_label(label):\n",
        "        if isinstance(label, str):\n",
        "            # Extract the first character (the number)\n",
        "            return int(label.split()[0])\n",
        "        else:\n",
        "            return int(label)\n",
        "\n",
        "    df[\"Label\"] = df[\"Label\"].apply(extract_label)\n",
        "    assert set(df[\"Label\"].unique()).issubset({1, 2, 3}), \"Labels must be 1/2/3.\"\n",
        "    return df\n",
        "\n",
        "train_df_synth = _ensure_int_labels(train_df)\n",
        "valid_df_food = _ensure_int_labels(valid_df)\n",
        "test_df_food = _ensure_int_labels(test_df)\n",
        "\n",
        "raw_dataset = DatasetDict(\n",
        "    {\n",
        "        \"train\": Dataset.from_pandas(train_df_synth),\n",
        "        \"validation\": Dataset.from_pandas(valid_df_food),\n",
        "        \"test\": Dataset.from_pandas(test_df_food),\n",
        "    }\n",
        ")\n",
        "\n",
        "print(f\"Train dataset size: {len(raw_dataset['train'])}\")\n",
        "print(f\"Validation dataset size: {len(raw_dataset['validation'])}\")\n",
        "print(f\"Test dataset size: {len(raw_dataset['test'])}\")\n",
        "print(f\"Train data sample: {raw_dataset['train'][0]}\")\n",
        "print(f\"Validation data sample: {raw_dataset['validation'][0]}\")\n",
        "print(f\"Test data sample {raw_dataset['test'][0]}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-tTsbiU1aYwt"
      },
      "source": [
        "To train our model, we need to convert our text prompts to token IDs. This is done by a Tokenizer, which tokenizes the inputs (including converting the tokens to their corresponding IDs in the pre-trained vocabulary) if you want to learn more about this, out **[chapter 6](https://huggingface.co/course/chapter6/1?fw=pt)** of the [Hugging Face Course](https://huggingface.co/course/chapter1/1)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "i5DyIB9UaYwu"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoTokenizer\n",
        "\n",
        "# Model id to load the tokenizer\n",
        "model_id = \"answerdotai/ModernBERT-base\"\n",
        "\n",
        "# Load Tokenizer\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
        "tokenizer.model_max_length = 512  # food names are short\n",
        "\n",
        "\n",
        "# Tokenize helper function\n",
        "def tokenize(batch):\n",
        "    return tokenizer(batch[\"Food\"], padding=\"max_length\", truncation=True)\n",
        "\n",
        "\n",
        "# Tokenize dataset\n",
        "# Keep original Label (1/2/3) for clarity, but create Trainer-compatible labels (0/1/2)\n",
        "raw_dataset = raw_dataset.map(\n",
        "    lambda b: {\"labels\": [int(x) - 1 for x in b[\"Label\"]]},\n",
        "    batched=True,\n",
        ")\n",
        "# If you want to drop raw text later, pass remove_columns=[\"Food\"]\n",
        "tokenized_dataset = raw_dataset.map(tokenize, batched=True)\n",
        "\n",
        "print(tokenized_dataset[\"train\"].features.keys())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s9Y0G0MJaYwu"
      },
      "source": [
        "### 2. Fine-tune & evaluate ModernBERT with the Hugging Face `Trainer`\n",
        "\n",
        "After we have processed our dataset, we can start training our model. We will use the [answerdotai/ModernBERT-base](https://huggingface.co/answerdotai/ModernBERT-base) model. The first step is to load our model with `AutoModelForSequenceClassification` class from the [Hugging Face Hub](https://huggingface.co/answerdotai/ModernBERT-base). This will initialize the pre-trained ModernBERT weights with a classification head on top. Here we pass the number of classes (3) from our dataset and the label names to have readable outputs for inference."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zFviwtalaYwv"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoModelForSequenceClassification\n",
        "import torch\n",
        "\n",
        "# Model id to load the tokenizer\n",
        "model_id = \"answerdotai/ModernBERT-base\"\n",
        "\n",
        "# Prepare model labels - useful for inference\n",
        "label2id = {name: str(i) for i, name in enumerate(label_names)}\n",
        "id2label = {str(i): name for i, name in enumerate(label_names)}\n",
        "\n",
        "# Download the model from huggingface.co/models\n",
        "model = AutoModelForSequenceClassification.from_pretrained(\n",
        "    model_id,\n",
        "    num_labels=num_labels,\n",
        "    label2id=label2id,\n",
        "    id2label=id2label,\n",
        "    attn_implementation=\"sdpa\",  # avoids flash-attn path\n",
        "    torch_dtype=torch.float16,    # good default on A100/L4/T4? (T4 prefers fp16)\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3VKhfQ6uaYww"
      },
      "source": [
        "We evaluate our model during training. The `Trainer` supports evaluation during training by providing a `compute_metrics` method. We use the `evaluate` library to calculate the [f1 metric](https://huggingface.co/spaces/evaluate-metric/f1) during training on our test split."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n3eKXvjlaYww"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "from sklearn.metrics import f1_score, accuracy_score\n",
        "\n",
        "# Metric helper method\n",
        "def compute_metrics(eval_pred):\n",
        "    predictions, labels = eval_pred\n",
        "    predictions = np.argmax(predictions, axis=1)\n",
        "\n",
        "    # Calculate accuracy\n",
        "    accuracy = accuracy_score(labels, predictions)\n",
        "\n",
        "    # Calculate weighted F1 score\n",
        "    f1 = f1_score(labels, predictions, average=\"weighted\")\n",
        "\n",
        "    return {\n",
        "        \"accuracy\": float(accuracy),\n",
        "        \"f1\": float(f1)\n",
        "    }"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r2FkowZ2aYww"
      },
      "source": [
        "The last step is to define the hyperparameters (`TrainingArguments`) we use for our training. Here we are adding optimizations introduced features for fast training times using `torch_compile` option in the `TrainingArguments`.\n",
        "\n",
        "We also leverage the [Hugging Face Hub](https://huggingface.co/models) integration of the `Trainer` to push our checkpoints, logs, and metrics during training into a repository."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AuALdtbyaYwx"
      },
      "outputs": [],
      "source": [
        "from huggingface_hub import HfFolder\n",
        "from transformers import Trainer, TrainingArguments\n",
        "\n",
        "# Define training args\n",
        "training_args = TrainingArguments(\n",
        "    output_dir= \"modernbert-llm-router\",\n",
        "    per_device_train_batch_size=32,\n",
        "    per_device_eval_batch_size=16,\n",
        "    learning_rate=5e-5,\n",
        "\t\tnum_train_epochs=1,\n",
        "    optim=\"adamw_torch_fused\", # improved optimizer\n",
        "    # logging & evaluation strategies\n",
        "    logging_strategy=\"steps\",\n",
        "    logging_steps=100,\n",
        "    eval_strategy=\"epoch\",\n",
        "    save_strategy=\"epoch\",\n",
        "    save_total_limit=2,\n",
        "    load_best_model_at_end=True,\n",
        "    metric_for_best_model=\"accuracy\",\n",
        "    # push to hub parameters\n",
        "    report_to=\"tensorboard\",\n",
        "\n",
        ")\n",
        "\n",
        "# Create a Trainer instance\n",
        "trainer = Trainer(\n",
        "    model=model,\n",
        "    args=training_args,\n",
        "    train_dataset=tokenized_dataset[\"train\"],\n",
        "    eval_dataset=tokenized_dataset[\"test\"],\n",
        "    compute_metrics=compute_metrics,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MRM2JiyhaYwx"
      },
      "source": [
        "We can start our training by using the **`train`** method of the `Trainer`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jnkTlEmtaYwx"
      },
      "outputs": [],
      "source": [
        "trainer.train()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-NjtKdOUgDrN"
      },
      "source": [
        "# Optional: Synthetic Dataset for Distillation\n",
        "\n",
        "Suppose you have a big, highly capable model, but you want to develop a much smaller model that has learned a set of skills from the big model.\n",
        "Oone way to do this is through model distillation--finetuning the small model to produce outputs that look like those of the big model.\n",
        "A very simple way to do model distillation is to generate a large number of examples using the big model, and then finetuning those examples on the small model.\n",
        "\n",
        "Below, we have provided you starter code that uses the OpenAI API to label 1,000 food items. As one approach for improving your BERT model, you may consider\n",
        "training on a larger data that of GPT-5-labeled foods."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s0gzEjg7yqZL"
      },
      "outputs": [],
      "source": [
        "# @title Load + sample unlabeled recipes\n",
        "!pip -q install datasets\n",
        "\n",
        "import json\n",
        "import re\n",
        "\n",
        "import pandas as pd\n",
        "from datasets import load_dataset\n",
        "\n",
        "DATASET_ID = \"Shengtao/recipe\"\n",
        "SEED = 42\n",
        "N_UNLABELED = 1000  # FOR STUDENT IMPL\n",
        "MAX_UNLABELED = 10000\n",
        "\n",
        "dataset = load_dataset(DATASET_ID, split=\"train\")  # FOR STUDENT IMPL\n",
        "raw_titles = dataset.shuffle(seed=SEED).select(range(min(len(dataset), MAX_UNLABELED * 3)))[\"title\"]\n",
        "\n",
        "s = (\n",
        "    pd.Series(raw_titles)\n",
        "    .astype(str)\n",
        "    .str.strip()\n",
        "    .str.replace(r\"\\s+(?=[IVXLCDM]+$)[IVXLCDM]+$\", \"\", regex=True)\n",
        ")\n",
        "clean_titles = s[s.ne(\"\") & ~s.str.lower().duplicated()].head(N_UNLABELED).tolist()\n",
        "\n",
        "unlabeled_df = pd.DataFrame({\"Food\": clean_titles})\n",
        "unlabeled_df.head()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qiALK6CZgvt8"
      },
      "outputs": [],
      "source": [
        "# @title Build Batch input JSONL for OpenAI\n",
        "from io import BytesIO\n",
        "from openai import OpenAI\n",
        "from google.colab import userdata\n",
        "\n",
        "client = OpenAI(api_key=userdata.get(\"OPENAI_API_KEY\"))  # FOR STUDENT IMPL\n",
        "\n",
        "SYSTEM_PROMPT = create_system_prompt(train_df)\n",
        "MODEL_NAME = \"gpt-5.2-2025-12-11\"  # FOR STUDENT IMPL\n",
        "\n",
        "\n",
        "def build_request(food_name, custom_id):\n",
        "    return {\n",
        "        \"custom_id\": custom_id,\n",
        "        \"method\": \"POST\",\n",
        "        \"url\": \"/v1/responses\",\n",
        "        \"body\": {\n",
        "            \"model\": MODEL_NAME,\n",
        "            \"input\": [\n",
        "                {\n",
        "                    \"role\": \"developer\",\n",
        "                    \"content\": [\n",
        "                        {\"type\": \"input_text\", \"text\": SYSTEM_PROMPT},\n",
        "                    ],\n",
        "                },\n",
        "                {\n",
        "                    \"role\": \"user\",\n",
        "                    \"content\": [\n",
        "                        {\"type\": \"input_text\", \"text\": food_name},\n",
        "                    ],\n",
        "                },\n",
        "            ],\n",
        "            \"text\": {\"format\": {\"type\": \"text\"}},\n",
        "            \"reasoning\": {\"effort\": \"low\", \"summary\": \"auto\"},\n",
        "        },\n",
        "    }\n",
        "\n",
        "\n",
        "batch_lines = [\n",
        "    json.dumps(build_request(food, custom_id=f\"food-{i}\"))\n",
        "    for i, food in enumerate(unlabeled_df[\"Food\"])\n",
        "]\n",
        "\n",
        "print(\"Batch input lines:\", len(batch_lines))\n",
        "display(unlabeled_df.head())\n",
        "display(pd.Series(batch_lines[:3], name=\"jsonl\"))\n",
        "\n",
        "batch_input_jsonl = \"\\n\".join(batch_lines)\n",
        "batch_bytes = batch_input_jsonl.encode(\"utf-8\")\n",
        "\n",
        "\n",
        "batch_file = BytesIO(batch_bytes)\n",
        "batch_file.name = \"openai_batch_input.jsonl\"  # required by OpenAI file upload"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PxSJ7FUCgvt8"
      },
      "outputs": [],
      "source": [
        "# @title Submit batch + poll status\n",
        "import time\n",
        "\n",
        "batch_input_file = client.files.create(\n",
        "    file=batch_file,\n",
        "    purpose=\"batch\",\n",
        ")  # FOR STUDENT IMPL\n",
        "\n",
        "print(\"Uploaded batch file id:\", batch_input_file.id)\n",
        "\n",
        "batch = client.batches.create(\n",
        "    input_file_id=batch_input_file.id,\n",
        "    endpoint=\"/v1/responses\",\n",
        "    completion_window=\"24h\",\n",
        "    metadata={\"job\": \"food-preference-synth\"},\n",
        ")  # FOR STUDENT IMPL\n",
        "\n",
        "print(\"Batch id:\", batch.id)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "53ILs5L0gvt8"
      },
      "source": [
        "You can check on the status of your batch at [this link](https://platform.openai.com/batches/batch_696c17cbec308190b210336da56d092a) or using the code block below.\n",
        "\n",
        "If your Colab runtime disconnects after submitting the batch, you do **not** need to re-submit. Just paste the `batch.id` into the `BATCH_ID` variable below.\n",
        "\n",
        "When the status below reads ``completed``, you are readt to head onto the next part."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MUOsfzXLgvt8"
      },
      "outputs": [],
      "source": [
        "# Batch prediction can take a while. Run this cell to see status at any time.\n",
        "# If your runtime restarted, paste your saved batch id here.\n",
        "BATCH_ID = \"\"  # e.g. \"batch_abc123\" (leave empty to use current batch variable)\n",
        "\n",
        "batch_id = BATCH_ID or batch.id\n",
        "batch = client.batches.retrieve(batch_id)  # FOR STUDENT IMPL\n",
        "counts = batch.request_counts\n",
        "print(\n",
        "    f\"Status: {batch.status} | total={counts.total} \"\n",
        "    f\"completed={counts.completed} failed={counts.failed}\"\n",
        ")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "v-NBZIxtgvt8"
      },
      "outputs": [],
      "source": [
        "assert batch.status == \"completed\", f\"Batch not completed: {batch.status}\"  # FOR STUDENT IMPL\n",
        "output_file_id = batch.output_file_id  # FOR STUDENT IMPL\n",
        "\n",
        "output_lines = client.files.content(output_file_id).read().decode(\"utf-8\").splitlines()  # FOR STUDENT IMPL\n",
        "print(\"Output lines:\", len(output_lines))\n",
        "\n",
        "\n",
        "def extract_label_from_response(resp_obj):\n",
        "    # response.output is a list; last item has content[0].text\n",
        "    output = resp_obj[\"response\"]['body'][\"output\"]\n",
        "    text = output[-1][\"content\"][0][\"text\"].strip()\n",
        "    label = int(text[0])\n",
        "    assert label in (1, 2, 3), f\"Unexpected label: {text}\"\n",
        "    return label\n",
        "\n",
        "\n",
        "records = []\n",
        "for line in output_lines:\n",
        "    obj = json.loads(line)\n",
        "    idx = int(obj[\"custom_id\"].split(\"-\")[1])\n",
        "    food = unlabeled_df.iloc[idx][\"Food\"]\n",
        "    label = extract_label_from_response(obj)\n",
        "    records.append({\"Food\": food, \"Label\": label})\n",
        "\n",
        "synth_df = pd.DataFrame.from_records(records)\n",
        "print(\"Labeled rows:\", len(synth_df))\n",
        "display(synth_df[\"Label\"].value_counts())\n",
        "display(synth_df.head(5))\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "include_colab_link": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
