HW4.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# CS 162 HW4\n",
        "\n",
        "<a href=\"https://colab.research.google.com/github/PlusLabNLP/cs162-hw4-w25/blob/main/HW4.ipynb\"><img alt=\"Colab Demo\" src=\"https://img.shields.io/badge/​-Open%20in%20Colab-blue?logo=googlecolab&logoColor=F9AB00&style=flat\"></a>\n",
        "\n",
        "In this assignment, you'll work with a recent Large Language Model [Gemma 2 2B](https://arxiv.org/pdf/2408.00118). You'll learn how to use the model and its tokenizer, generate text using greedy decoding, top-p sampling, and top-k sampling, and evaluate the model’s basic arithmetic capabilities on a simple dataset.\n",
        "\n",
        "* Identify the TODO blocks and implement the necessary code in those sections.\n",
        "*  To speed up processing, use a GPU by selecting \"Runtime\" > \"Change runtime type\" > \"GPU\" in Colab."
      ],
      "metadata": {
        "id": "R5fMJ8rDNgQR"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 0 HuggingFace"
      ],
      "metadata": {
        "id": "DiTKWoDgMF8O"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m6dM_GtXLuPK"
      },
      "outputs": [],
      "source": [
        "# TODO: Generate a new Access Token (Read) at https://huggingface.co/settings/tokens\n",
        "# and enter it below to enable access to Gemma models, which are gated.\n",
        "# Ensure you have requested access to Gemma models at\n",
        "# https://huggingface.co/google/gemma-2-2b and received approval before proceeding.\n",
        "\n",
        "from huggingface_hub import notebook_login\n",
        "notebook_login()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3lxQAU8VLuPM"
      },
      "source": [
        "# 1 Model Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_NQeW-MSLuPN"
      },
      "outputs": [],
      "source": [
        "# For auto grading: Do not modify\n",
        "import torch\n",
        "\n",
        "def grade_model_setup(model):\n",
        "    ## model parameter type\n",
        "    first_param = next(model.parameters())\n",
        "    print(f\"Model parameter dtype: {first_param.dtype}\")\n",
        "\n",
        "    ## which device the model is on\n",
        "    device_idx = next(model.parameters()).get_device()\n",
        "    device = torch.cuda.get_device_name(device_idx) if device_idx != -1 else \"CPU\"\n",
        "    print(f\"Model is currently on device: {device}\")\n",
        "\n",
        "    ## what is the memory footprint\n",
        "    print(\"Memory:\", model.get_memory_footprint())\n",
        "\n",
        "    if first_param.dtype == torch.bfloat16:\n",
        "        print(\"<<Passed 1.1>>\")\n",
        "    else:\n",
        "        raise(Exception(f\"Failed 1.1: dtype is {first_param.dtype} instead of torch.bfloat16\"))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WYXkcSQiLuPO"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
        "\n",
        "\n",
        "def model_and_tokenizer_setup(model_name: str):\n",
        "    model, tokenizer = None, None\n",
        "    ##################################################\n",
        "    # TODO: Please load the model and tokenizer, which will\n",
        "    # be later used for inference. To have an optimized\n",
        "    # version of the model, load it in bfloat16 using torch_dtype.\n",
        "    # use AutoTokenizer and AutoModelForCausalLM.\n",
        "    # Hint: https://huggingface.co/google/gemma-2-2b#running-the-model-on-a-single--multi-gpu\n",
        "    pass\n",
        "    # End of TODO.\n",
        "    ##################################################\n",
        "    model.eval()\n",
        "    model.to('cuda')\n",
        "    grade_model_setup(model)\n",
        "    return model, tokenizer\n",
        "\n",
        "model, tokenizer = model_and_tokenizer_setup(model_name=\"google/gemma-2-2b\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OdEP6K7QLuPO"
      },
      "outputs": [],
      "source": [
        "# Do not modify\n",
        "def generate_next_token(model, tokenizer, tokenized_input, verbose=True):\n",
        "    \"\"\"\n",
        "    Generate the probability distribution over vocabulary\n",
        "    for the next token after the tokenized_input using the model.\n",
        "    \"\"\"\n",
        "    outputs = model.generate(**tokenized_input, max_new_tokens=1, output_scores=True, return_dict_in_generate=True)\n",
        "    next_token_scores = outputs[\"scores\"][0].squeeze()\n",
        "    next_token_scores = torch.softmax(next_token_scores, dim=0)\n",
        "    next_token_id_greedy = outputs[\"sequences\"][0][-1]\n",
        "    if verbose:\n",
        "        print(f\"tokenized_input: {tokenized_input}\")\n",
        "        print(f\"Shape of next_token_scores: {next_token_scores.shape}\\nnext_token_scores[:5]: {next_token_scores[:5]}\")\n",
        "    return next_token_scores\n",
        "\n",
        "input_text = \"Go Bruins\"\n",
        "tokenized_input = tokenizer(input_text, return_tensors=\"pt\").to(\"cuda\")\n",
        "next_token_scores = generate_next_token(model, tokenizer, tokenized_input, verbose=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tjPvnc0CLuPP"
      },
      "source": [
        "In transformer-based models like those from Hugging Face's transformers library, tokenized inputs are typically in batch format.\n",
        "\n",
        "* input_ids: This represents the tokenized numerical IDs of the input text. Since models process inputs in batches, tokenized_input['input_ids'] is a 2D list (batch_size × sequence_length). ['input_ids'][0] extracts the first example in the batch.\n",
        "\n",
        "* attention_mask: This tells the model which tokens to attend to (1) and which to ignore (0, usually for padding). It has the same shape as input_ids and ensures that padding tokens don’t affect the model’s output.\n",
        "\n",
        "* Since models process inputs as batches, even when working with a single sentence, the data is structured as a batch of size 1.\n",
        "\n",
        "* next_token_scores: A 1D tensor (torch.Size([256000])) containing the probabilities for each possible next token in the vocabulary after applying the softmax function. The size of 256,000 corresponds to the vocabulary size, meaning the model is considering 256,000 possible tokens for the next step in the sequence."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_xeiOVj0LuPP"
      },
      "outputs": [],
      "source": [
        "def decode_token_ids(tokenizer, token_ids: list) -> str:\n",
        "    decoded_text = None\n",
        "    ##################################################\n",
        "    # TODO: Please decode a list of token_ids to a string using the tokenizer.decode.\n",
        "    pass\n",
        "    # End of TODO.\n",
        "    ##################################################\n",
        "    return decoded_text\n",
        "\n",
        "\n",
        "decoded_tokens = decode_token_ids(tokenizer, tokenized_input['input_ids'][0])\n",
        "print(f\"Decoded output: {decoded_tokens}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "F9NzFdZBLuPQ"
      },
      "outputs": [],
      "source": [
        "# For auto grading: Do not modify\n",
        "def grade_tokenizer_setup():\n",
        "    text = decode_token_ids(tokenizer, [2, 6196, 13435, 575, 6461, 603, 78672])\n",
        "    correct_text = \"<bos>Best university in LA is UCLA\"\n",
        "    print(text)\n",
        "    if text == correct_text:\n",
        "        print(f\"<<Passed 1.2>>\")\n",
        "    else:\n",
        "        raise(Exception(f\"Failed 1.2: decoded text is '{text}' instead of '{correct_text}'\"))\n",
        "\n",
        "grade_tokenizer_setup()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WlJVqgOoLuPQ"
      },
      "source": [
        "# 2 Decoding Strategies\n",
        "In this section, you should implement different decoding strategies for text generation:\n",
        "1. Greedy Decoding: Always selects the highest probability token at each step.\n",
        "2. Top-P Sampling (Nucleus Sampling): Selects from a subset of the vocabulary containing the top-P cumulative probability mass.\n",
        "3. Top-K Sampling: Selects from the top-K most probable tokens at each step."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rh_gGDhdLuPQ"
      },
      "outputs": [],
      "source": [
        "from tqdm.auto import tqdm\n",
        "\n",
        "def append_next_token(tokenized_input, next_token_id):\n",
        "    \"\"\"\n",
        "    Appends the next token to the tokenized input sequence.\n",
        "\n",
        "    This function is used to incrementally build input sequences for autoregressive decoding.\n",
        "    It takes a tokenized input (the result of running a tokenizer on text) and a next_token_id\n",
        "    (the ID of the next token predicted by a model) and appends this token ID to the input sequence.\n",
        "    It also updates the attention mask accordingly.\n",
        "\n",
        "    Args:\n",
        "        tokenized_input: A dictionary containing:\n",
        "            - \"input_ids\": A tensor of token IDs representing the input sequence.\n",
        "            - \"attention_mask\": A tensor indicating which tokens should be attended to (1 for real tokens, 0 for padding).\n",
        "        next_token_id: The ID of the next token to append, predicted by a model.\n",
        "\n",
        "    Returns:\n",
        "        Updated tokenized input with the new token ID and attention mask.\n",
        "    \"\"\"\n",
        "    tokenized_input[\"input_ids\"] = torch.cat(\n",
        "        [tokenized_input[\"input_ids\"], torch.tensor([[next_token_id]], dtype=torch.int64).to(\"cuda\")], dim=-1\n",
        "    )\n",
        "    tokenized_input[\"attention_mask\"] = torch.cat(\n",
        "        [tokenized_input[\"attention_mask\"], torch.tensor([[1]], dtype=torch.int64).to(\"cuda\")], dim=-1\n",
        "    )\n",
        "    return tokenized_input\n",
        "\n",
        "\n",
        "def greedy_decoding(model, tokenizer, input_text, max_length, disable_tqdm=True):\n",
        "    \"\"\"\n",
        "    Generates text using greedy decoding by always selecting the most probable next token.\n",
        "\n",
        "    Args:\n",
        "        model: The language model used for text generation.\n",
        "        tokenizer: The tokenizer corresponding to the model.\n",
        "        input_text: The initial text prompt to generate from.\n",
        "        max_length: The maximum number of new tokens to generate.\n",
        "\n",
        "    Returns:\n",
        "        sequence_str: The whole text as a string (input and generated).\n",
        "    \"\"\"\n",
        "    sequence = tokenizer(input_text, return_tensors=\"pt\").to(\"cuda\")\n",
        "    for _ in tqdm(range(max_length), disable=disable_tqdm):\n",
        "        next_token_scores = generate_next_token(model, tokenizer, sequence, verbose=False)\n",
        "        ##################################################\n",
        "        # TODO: Implement greedy decoding using the next_token_scores.\n",
        "        # Hint: generate tokens one by one until max_length is reached or the model generates <eos> (tokenizer.eos_token_id).\n",
        "        # Hint: \"sequence\" stores the tokenized input text and is updated with newly generated tokens containing the whole input and new generated tokens.\n",
        "        # Hint: Use append_next_token function above to update \"sequence\" with the new token you generate.\n",
        "        # Hint: Use torch.argmax to get the most probable token.\n",
        "        pass\n",
        "        # End of TODO.\n",
        "        ##################################################\n",
        "\n",
        "    sequence_str = None\n",
        "    ##################################################\n",
        "    # TODO: Use decode_token_ids to decode the sequence to text.\n",
        "    pass\n",
        "    # End of TODO.\n",
        "    ##################################################\n",
        "    return sequence_str\n",
        "\n",
        "def top_p_sampling(model, tokenizer, input_text, max_length, top_p, seed, disable_tqdm=True):\n",
        "    \"\"\"\n",
        "    Generates text using top-p (nucleus) sampling by sampling from the smallest subset\n",
        "    of tokens whose cumulative probability mass exceeds `top_p`.\n",
        "\n",
        "    Args:\n",
        "        model: The language model used for text generation.\n",
        "        tokenizer: The tokenizer corresponding to the model.\n",
        "        input_text: The initial text prompt to generate from.\n",
        "        max_length: The maximum number of new tokens to generate.\n",
        "        top_p: The nucleus sampling threshold.\n",
        "\n",
        "    Returns:\n",
        "        sequence_str: The whole text as a string (input and generated).\n",
        "    \"\"\"\n",
        "    torch.random.manual_seed(seed)\n",
        "    sequence_str = None\n",
        "    ##################################################\n",
        "    # TODO: Implement top-p sampling using the generate_next_token function.\n",
        "    # Hint: Pay attention to greedy_decoding function and how it generates tokens.\n",
        "    # Hint: Use torch argsort, cumsum, and multinomial to sample from the top-p tokens.\n",
        "    pass\n",
        "    # End of TODO.\n",
        "    ##################################################\n",
        "    return sequence_str\n",
        "\n",
        "def top_k_sampling(model, tokenizer, input_text, max_length, top_k, seed, disable_tqdm=True):\n",
        "    \"\"\"\n",
        "    Generates text using top-k sampling by selecting from the top-K most probable tokens.\n",
        "\n",
        "    Args:\n",
        "        model: The language model used for text generation.\n",
        "        tokenizer: The tokenizer corresponding to the model.\n",
        "        input_text: The initial text prompt to generate from.\n",
        "        max_length: The maximum number of new tokens to generate.\n",
        "        top_k: The number of top candidates to sample from.\n",
        "\n",
        "    Returns:\n",
        "        sequence_str: The whole text as a string (input and generated).\n",
        "    \"\"\"\n",
        "    torch.random.manual_seed(seed)\n",
        "    sequence_str = None\n",
        "    ##################################################\n",
        "    # TODO: Implement top-k sampling using the generate_next_token function.\n",
        "    # Hint: Pay attention to greedy_decoding function and how it generates tokens.\n",
        "    # Hint: Use torch topk and multinomial to sample from the top-k tokens.\n",
        "    pass\n",
        "    # End of TODO.\n",
        "    ##################################################\n",
        "    return sequence_str\n",
        "\n",
        "print(\"greedy_decoding:\", greedy_decoding(model, tokenizer, input_text=\"Best university in LA is \", max_length=45, disable_tqdm=False))\n",
        "print(\"top_p_sampling:\", top_p_sampling(model, tokenizer, input_text=\"Best university in LA is \", max_length=20, top_p=0.8, seed=0, disable_tqdm=False))\n",
        "print(\"top_k_sampling:\", top_k_sampling(model, tokenizer, input_text=\"Best university in LA is \", max_length=20, top_k=200, seed=0, disable_tqdm=False))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zsyosp5MLuPR"
      },
      "outputs": [],
      "source": [
        "# For auto grading: Do not modify\n",
        "\n",
        "def grade_greedy_decoding():\n",
        "    r = greedy_decoding(model, tokenizer, input_text=\"Best university in LA is \", max_length=4)\n",
        "    correct_text = \"<bos>Best university in LA is <strong>UCLA</strong>.\"\n",
        "    if r == correct_text:\n",
        "        print(f\"<<Passed 2.1>>\")\n",
        "    else:\n",
        "        print(f\"###Failed 2.1###: generated text is '{r}' instead of '{correct_text}'\")\n",
        "\n",
        "def grade_top_p_sampling():\n",
        "    r = top_p_sampling(model, tokenizer, input_text=\"Best university in LA is \", max_length=4, top_p=0.0, seed=3)\n",
        "    correct_text = \"<bos>Best university in LA is <strong>UCLA</strong>.\"\n",
        "    if r == correct_text:\n",
        "        print(f\"<<Passed 2.2.1>>\")\n",
        "    else:\n",
        "        print(f\"###Failed 2.2.1###: generated text is '{r}' instead of '{correct_text}'\")\n",
        "\n",
        "def grade_top_k_sampling():\n",
        "    r = top_k_sampling(model, tokenizer, input_text=\"Best university in LA is \", max_length=4, top_k=1, seed=3)\n",
        "    correct_text = \"<bos>Best university in LA is <strong>UCLA</strong>.\"\n",
        "    if r == correct_text:\n",
        "        print(f\"<<Passed 2.3.1>>\")\n",
        "    else:\n",
        "        print(f\"###Failed 2.3.1###: generated text is '{r}' instead of '{correct_text}'\")\n",
        "\n",
        "grade_greedy_decoding()\n",
        "grade_top_p_sampling()\n",
        "grade_top_k_sampling()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ssc7ejtmLuPR"
      },
      "outputs": [],
      "source": [
        "NUM_RUNS = 3\n",
        "MAX_LENGTH = 20\n",
        "TOP_P = 0.8\n",
        "TOP_K = 10\n",
        "INPUT_TEXT = \"My new AI assistant just claimed it invented \"\n",
        "##################################################\n",
        "# Please explore different decoding strategies for language generation by running three algorithms\n",
        "# —greedy decoding, top-p (nucleus) sampling, and top-k sampling—multiple times and printing their outputs.\n",
        "# Specifically, you should use the functions `greedy_decoding`, `top_p_sampling`, and `top_k_sampling` to generate text.\n",
        "# For each algorithm, run it using ***different seeds*** (you can use i from the for loop) and observe how the outputs vary across runs.\n",
        "# Since greedy decoding is deterministic, its outputs should remain consistent, while top-p and top-k sampling introduce randomness, leading to different generations.\n",
        "# Set the values for top-p and top-k as TOP_P and TOP_K.\n",
        "# Use the `max_length` of MAX_LENGTH for all generations to ensure consistency in output length.\n",
        "# Your goal is to understand these behaviors by observing the diversity in the generated outputs.\n",
        "# When submitting your notebook, make sure to include the outputs of your runs for manual grading. (Do not clear outputs)\n",
        "\n",
        "# GREEDY DECODING\n",
        "print(\"######### Greedy Decoding: #########\")\n",
        "for i in range(NUM_RUNS):\n",
        "    ##################################################\n",
        "    # TODO: Please run greedy decoding and print the output.\n",
        "    pass\n",
        "    # End of TODO.\n",
        "    ##################################################\n",
        "\n",
        "print(\"######### Top-p Sampling: #########\")\n",
        "for i in range(NUM_RUNS):\n",
        "    ##################################################\n",
        "    # TODO: Please run top-p sampling and print the output.\n",
        "    pass\n",
        "    # End of TODO.\n",
        "    ##################################################\n",
        "\n",
        "print(\"######### Top-k Sampling: #########\")\n",
        "for i in range(NUM_RUNS):\n",
        "    ##################################################\n",
        "    # TODO: Please run top-k sampling and print the output.\n",
        "    pass\n",
        "    # End of TODO.\n",
        "    ##################################################"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sbq2_izhLuPR"
      },
      "source": [
        "# 3 Load Arithmetic Dataset\n",
        "\n",
        "In this section, we will implement a data loader using torch.utils.data.Dataset to load a file containing arithmetic problems and corresponding answers."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Rha8XF0yLuPR"
      },
      "outputs": [],
      "source": [
        "import pandas as pd\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "\n",
        "\n",
        "class ArithmeticDataset(Dataset):\n",
        "    \"\"\"\n",
        "    A custom dataset class for loading arithmetic problems and answers from a file.\n",
        "    The dataset can be filtered based on operation type and number of digits in the operands.\n",
        "    \"\"\"\n",
        "    def __init__(self, file_path_csv, operation=None, num_digits=None):\n",
        "        \"\"\"\n",
        "        Initialize the dataset with the file containing arithmetic problems and answers.\n",
        "\n",
        "        Args:\n",
        "        - file_path_csv: Path to the csv file containing the arithmetic problems and their answers.\n",
        "        - operation: Optional filter for the arithmetic operation. Can be 'add', 'sub', 'mul', or 'div'.\n",
        "        - num_digits: Optional filter for the number of digits in the operands.\n",
        "        \"\"\"\n",
        "        self.df = pd.read_csv(file_path_csv)\n",
        "        self.operation = operation\n",
        "        self.num_digits = num_digits\n",
        "\n",
        "        ##################################################\n",
        "        # TODO: Please complete the implementation of __init__\n",
        "        # filter the dataframe (self.df) as per the operation type and num_digits given to you.\n",
        "        # (Do not filter operation or num_digits if they are None.)\n",
        "        pass\n",
        "        # End of TODO.\n",
        "        ##################################################\n",
        "\n",
        "    def __len__(self):\n",
        "        \"\"\"Return the number of examples in the dataset.\"\"\"\n",
        "        return len(self.df)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        \"\"\"\n",
        "        Get an item (problem, answer) by index.\n",
        "\n",
        "        Args:\n",
        "        - idx: Index of the item to retrieve.\n",
        "\n",
        "        Returns:\n",
        "        - A tuple containing the problem (as a string) and the corresponding answer (as an integer).\n",
        "        \"\"\"\n",
        "        problem, answer = None, None\n",
        "        ##################################################\n",
        "        # TODO: Please complete the implementation of __getitem__\n",
        "        # to return the problem and answer at the given index.\n",
        "        # Hint: Use the df.iloc to access the index.\n",
        "        pass\n",
        "        # End of TODO.\n",
        "        ##################################################\n",
        "        return problem, answer\n",
        "\n",
        "dataset = ArithmeticDataset(file_path_csv=\"arithmetic_problems.csv\")\n",
        "dataset.df"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N8g9cIigLuPS"
      },
      "outputs": [],
      "source": [
        "# For auto grading: Do not modify\n",
        "def grade_dataset():\n",
        "    dataset = ArithmeticDataset(file_path_csv=\"arithmetic_problems.csv\", operation=\"add\", num_digits=2)\n",
        "    if len(dataset) != 14:\n",
        "        raise(Exception(f\"Failed 3.1: dataset length is {len(dataset)} instead of 14\"))\n",
        "    if dataset[10] != ('62 + 31 = ', 93):\n",
        "        raise(Exception(f\"Failed 3.1: dataset[10] is {dataset[10]} instead of ('62 + 31 = ', 93)\"))\n",
        "    print(\"<<Passed 3.1>>\")\n",
        "\n",
        "grade_dataset()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y6nJS0STLuPS"
      },
      "source": [
        "# 4 Implement Evaluation Metric"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TSIP7mwQLuPS"
      },
      "outputs": [],
      "source": [
        "from sklearn.metrics import accuracy_score, f1_score\n",
        "\n",
        "def evaluate_standard(gt_answers, pred_answers):\n",
        "\n",
        "    accuracy = 0\n",
        "\n",
        "    ##################################################\n",
        "    # TODO: Please finish the standard evaluation metrics.\n",
        "    # You need to compute the accuracy for the\n",
        "    # predictions and ground truth answers.\n",
        "    # Please use the scikit-learn APIs.\n",
        "\n",
        "    pass\n",
        "\n",
        "    # End of TODO.\n",
        "    ##################################################\n",
        "\n",
        "    return accuracy"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "943bu_3fLuPS"
      },
      "outputs": [],
      "source": [
        "# For auto grading: Do not modify\n",
        "def grade_evaluation_metric():\n",
        "    gt_answers = [1, 2, 3, 4, 5]\n",
        "    pred_answers = [1, 2, 3, 4, 2]\n",
        "    acc = evaluate_standard(gt_answers, pred_answers)\n",
        "    if acc != 0.8:\n",
        "        raise(Exception(f\"Failed 4.1: accuracy is {acc} instead of 0.8\"))\n",
        "    print(\"<<Passed 4.1>>\")\n",
        "\n",
        "grade_evaluation_metric()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "klc4TCtxLuPS"
      },
      "source": [
        "# 5 Evaluate the Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "R0woTLA4LuPS"
      },
      "outputs": [],
      "source": [
        "def evaluate_arithmetic(model, tokenizer, dataset_path_csv, decoding_strategy, top_p=None, top_k=None):\n",
        "    \"\"\"\n",
        "    Evaluate the model on the arithmetic dataset using the given decoding strategy.\n",
        "\n",
        "    Args:\n",
        "    - model: The language model used for text generation.\n",
        "    - tokenizer: The tokenizer corresponding to the model.\n",
        "    - dataset_path_csv: The path to the csv file containing the arithmetic problems and their answers.\n",
        "    - decoding_strategy: The decoding strategy to use for text generation. (Can be \"greedy\", \"top_p\", \"top_k\")\n",
        "    - top_p: The nucleus (top-p) sampling threshold. (Only used if decoding_strategy is \"top_p\")\n",
        "    - top_k: The number of top candidates to sample from. (Only used if decoding_strategy is \"top_k\")\n",
        "\n",
        "    Returns:\n",
        "    - answers: A pandas DataFrame containing the evaluation results.\n",
        "    \"\"\"\n",
        "    answers = []\n",
        "    for operation in tqdm([\"add\", \"sub\", \"mul\", \"div\"]):\n",
        "        for num_digits in [1, 2, 3, 4]:\n",
        "            ##################################################\n",
        "            # TODO: Please load the dataset for the given operation and num_digits.\n",
        "            # Hint: Use the ArithmeticDataset class you implemented.\n",
        "            # Hint: Use the dataset_path_csv, operation, and num_digits arguments.\n",
        "            dataset = None\n",
        "            # End of TODO.\n",
        "            ##################################################\n",
        "\n",
        "            for problem, answer in DataLoader(dataset, batch_size=1):\n",
        "                pred_sequence = None\n",
        "                ##################################################\n",
        "                # TODO: Generate the prediction for the given problem based on the decoding_strategy.\n",
        "                # Hint: Use the greedy_decoding, top_k_sampling, and top_p_sampling functions.\n",
        "                # Hint: You can use if statements to call the appropriate function based on the decoding_strategy.\n",
        "                # IMPORTANT: For max_length, use the length of the answer (How many digits are in the answer).\n",
        "                # Hint: answer is a tensor, so you can access the value itself using answer.item() method and then count its number of digits.\n",
        "                # For seed, use 0.\n",
        "                pass\n",
        "                # End of TODO.\n",
        "                ##################################################\n",
        "\n",
        "                # Extracting final answer from the sequence (What comes after \"=\")\n",
        "                pred_answer = pred_sequence.split(\"=\")[1].strip()\n",
        "                answers.append({\n",
        "                    \"operation\": operation,\n",
        "                    \"num_digits\": num_digits,\n",
        "                    \"problem\": problem[0],\n",
        "                    \"sequence\": pred_sequence,\n",
        "                    \"decoding_strategy\": decoding_strategy,\n",
        "                    \"true_answer\": str(answer.item()),\n",
        "                    \"pred_answer\": pred_answer\n",
        "                })\n",
        "    return pd.DataFrame(answers)\n",
        "\n",
        "answers = evaluate_arithmetic(model, tokenizer, \"arithmetic_problems.csv\", \"greedy\")\n",
        "answers"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3V7nKG4wLuPT"
      },
      "outputs": [],
      "source": [
        "# See the incorrect predictions\n",
        "answers[answers[\"true_answer\"] != answers[\"pred_answer\"]]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Uxnn8AkDLuPT"
      },
      "outputs": [],
      "source": [
        "accuracy = None\n",
        "##################################################\n",
        "# TODO: Evaluate the predictions using the evaluate_standard function.\n",
        "# Hint: Use the \"true_answer\" and \"pred_answer\" columns from the answers dataframe.\n",
        "pass\n",
        "# End of TODO.\n",
        "##################################################\n",
        "accuracy"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aIncyFqLLuPT"
      },
      "outputs": [],
      "source": [
        "# For auto grading: Do not modify\n",
        "def grade_acc():\n",
        "    if 0.75 < accuracy < 1:\n",
        "        print(\"<<Passed 5.1>>\")\n",
        "    else:\n",
        "        raise(Exception(f\"Failed 5.1: accuracy is {accuracy} instead of 0.75<acc<1\"))\n",
        "\n",
        "grade_acc()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "byRbmI_mLuPT"
      },
      "outputs": [],
      "source": [
        "# For auto grading: Do not modify\n",
        "def grade_operation_and_digits():\n",
        "    dfs = {}\n",
        "    dfs[\"greedy\"] = evaluate_arithmetic(model, tokenizer, \"arithmetic_problems.csv\", \"greedy\")\n",
        "    dfs[\"top_p\"] = evaluate_arithmetic(model, tokenizer, \"arithmetic_problems.csv\", \"top_p\", top_p=0.9)\n",
        "    dfs[\"top_k\"] = evaluate_arithmetic(model, tokenizer, \"arithmetic_problems.csv\", \"top_k\", top_k=20)\n",
        "    for k in dfs.keys():\n",
        "        dfs[k][\"acc\"] = dfs[k][\"true_answer\"] == dfs[k][\"pred_answer\"]\n",
        "        dfs[k] = dfs[k].groupby([\"operation\", \"num_digits\"])[\"acc\"].mean().reset_index()\n",
        "    df = dfs[\"greedy\"].copy()[[\"operation\", \"num_digits\"]]\n",
        "    df[\"acc_greedy\"] = dfs[\"greedy\"][\"acc\"] * 100\n",
        "    df[\"acc_top_p\"] = dfs[\"top_p\"][\"acc\"] * 100\n",
        "    df[\"acc_top_k\"] = dfs[\"top_k\"][\"acc\"] * 100\n",
        "    df = df.sort_values([\"num_digits\", \"operation\"]).reset_index(drop=True)\n",
        "    df.loc['mean'] = df[[\"acc_greedy\", \"acc_top_p\", \"acc_top_k\"]].mean()\n",
        "\n",
        "    def check_acc(name, value, valid_range):\n",
        "        if valid_range[0] < value < valid_range[1]:\n",
        "            print(f\"<<Passed 5.2.{name}>>\")\n",
        "        else:\n",
        "            print(Exception(f\"Failed 5.2.{name}: {value} is not in the range {valid_range}\"))\n",
        "\n",
        "    check_acc(\"acc_greedy\", df.loc['mean', 'acc_greedy'], (75, 100))\n",
        "    check_acc(\"acc_top_p\", df.loc['mean', 'acc_top_p'], (50, 70))\n",
        "    check_acc(\"acc_top_k\", df.loc['mean', 'acc_top_k'], (50, 70))\n",
        "    return df.style.background_gradient(cmap='RdYlGn', subset=[\"acc_greedy\", \"acc_top_p\", \"acc_top_k\"]).format(precision=2)\n",
        "\n",
        "grade_operation_and_digits()"
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "4IDKaJYsSNll"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "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.5"
    },
    "colab": {
      "provenance": [],
      "toc_visible": true,
      "gpuType": "T4"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}