{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# CPT Training and Inference\n",
        "This notebook demonstrates the training and evaluation process of Context-Aware Prompt Tuning (CPT) using the Hugging Face Trainer. For more details, refer to the [Paper](https://huggingface.co/papers/2410.17222).\n",
        "\n",
        "\n",
        "## Sections Overview:\n",
        "1. **Setup**: Import libraries and configure the environment.\n",
        "2. **Data Preparation**: Load and preprocess the dataset.\n",
        "3. **Model Training**: Configure and train the model.\n",
        "4. **Evaluation**: Test the model's performance and visualize results."
      ],
      "metadata": {
        "id": "R_byvXT9lpTU"
      },
      "id": "R_byvXT9lpTU"
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Setup\n",
        "\n",
        "---\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "collapsed": false,
        "id": "11b07b07ac5e472b"
      },
      "id": "11b07b07ac5e472b"
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Installation"
      ],
      "metadata": {
        "id": "O8DWZb8ZrGRU"
      },
      "id": "O8DWZb8ZrGRU"
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install datasets\n",
        "!pip install git+https://github.com/huggingface/peft"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "d6KZ5REDrFiM",
        "outputId": "e505bc0e-082a-4720-9117-b730d9fd67fa"
      },
      "id": "d6KZ5REDrFiM",
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: datasets in /usr/local/lib/python3.10/dist-packages (3.1.0)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from datasets) (3.16.1)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from datasets) (1.26.4)\n",
            "Requirement already satisfied: pyarrow>=15.0.0 in /usr/local/lib/python3.10/dist-packages (from datasets) (17.0.0)\n",
            "Requirement already satisfied: dill<0.3.9,>=0.3.0 in /usr/local/lib/python3.10/dist-packages (from datasets) (0.3.8)\n",
            "Requirement already satisfied: pandas in /usr/local/lib/python3.10/dist-packages (from datasets) (2.2.2)\n",
            "Requirement already satisfied: requests>=2.32.2 in /usr/local/lib/python3.10/dist-packages (from datasets) (2.32.3)\n",
            "Requirement already satisfied: tqdm>=4.66.3 in /usr/local/lib/python3.10/dist-packages (from datasets) (4.66.6)\n",
            "Requirement already satisfied: xxhash in /usr/local/lib/python3.10/dist-packages (from datasets) (3.5.0)\n",
            "Requirement already satisfied: multiprocess<0.70.17 in /usr/local/lib/python3.10/dist-packages (from datasets) (0.70.16)\n",
            "Requirement already satisfied: fsspec<=2024.9.0,>=2023.1.0 in /usr/local/lib/python3.10/dist-packages (from fsspec[http]<=2024.9.0,>=2023.1.0->datasets) (2024.9.0)\n",
            "Requirement already satisfied: aiohttp in /usr/local/lib/python3.10/dist-packages (from datasets) (3.11.2)\n",
            "Requirement already satisfied: huggingface-hub>=0.23.0 in /usr/local/lib/python3.10/dist-packages (from datasets) (0.26.2)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from datasets) (24.2)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from datasets) (6.0.2)\n",
            "Requirement already satisfied: aiohappyeyeballs>=2.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (2.4.3)\n",
            "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (1.3.1)\n",
            "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (24.2.0)\n",
            "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (1.5.0)\n",
            "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (6.1.0)\n",
            "Requirement already satisfied: propcache>=0.2.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (0.2.0)\n",
            "Requirement already satisfied: yarl<2.0,>=1.17.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (1.17.2)\n",
            "Requirement already satisfied: async-timeout<6.0,>=4.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (4.0.3)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.23.0->datasets) (4.12.2)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests>=2.32.2->datasets) (3.4.0)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests>=2.32.2->datasets) (3.10)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests>=2.32.2->datasets) (2.2.3)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests>=2.32.2->datasets) (2024.8.30)\n",
            "Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.10/dist-packages (from pandas->datasets) (2.8.2)\n",
            "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas->datasets) (2024.2)\n",
            "Requirement already satisfied: tzdata>=2022.7 in /usr/local/lib/python3.10/dist-packages (from pandas->datasets) (2024.2)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil>=2.8.2->pandas->datasets) (1.16.0)\n",
            "Collecting git+https://github.com/huggingface/peft\n",
            "  Cloning https://github.com/huggingface/peft to /tmp/pip-req-build-0mbyx_z_\n",
            "  Running command git clone --filter=blob:none --quiet https://github.com/huggingface/peft /tmp/pip-req-build-0mbyx_z_\n",
            "  Resolved https://github.com/huggingface/peft to commit 131efba5d48753a3355ecd4f3833ae010a0510d6\n",
            "  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from peft==0.13.3.dev0) (1.26.4)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from peft==0.13.3.dev0) (24.2)\n",
            "Requirement already satisfied: psutil in /usr/local/lib/python3.10/dist-packages (from peft==0.13.3.dev0) (5.9.5)\n",
            "Requirement already satisfied: pyyaml in /usr/local/lib/python3.10/dist-packages (from peft==0.13.3.dev0) (6.0.2)\n",
            "Requirement already satisfied: torch>=1.13.0 in /usr/local/lib/python3.10/dist-packages (from peft==0.13.3.dev0) (2.5.1+cu121)\n",
            "Requirement already satisfied: transformers in /usr/local/lib/python3.10/dist-packages (from peft==0.13.3.dev0) (4.46.2)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from peft==0.13.3.dev0) (4.66.6)\n",
            "Requirement already satisfied: accelerate>=0.21.0 in /usr/local/lib/python3.10/dist-packages (from peft==0.13.3.dev0) (1.1.1)\n",
            "Requirement already satisfied: safetensors in /usr/local/lib/python3.10/dist-packages (from peft==0.13.3.dev0) (0.4.5)\n",
            "Requirement already satisfied: huggingface_hub>=0.25.0 in /usr/local/lib/python3.10/dist-packages (from peft==0.13.3.dev0) (0.26.2)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from huggingface_hub>=0.25.0->peft==0.13.3.dev0) (3.16.1)\n",
            "Requirement already satisfied: fsspec>=2023.5.0 in /usr/local/lib/python3.10/dist-packages (from huggingface_hub>=0.25.0->peft==0.13.3.dev0) (2024.9.0)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from huggingface_hub>=0.25.0->peft==0.13.3.dev0) (2.32.3)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface_hub>=0.25.0->peft==0.13.3.dev0) (4.12.2)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.0->peft==0.13.3.dev0) (3.4.2)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.0->peft==0.13.3.dev0) (3.1.4)\n",
            "Requirement already satisfied: sympy==1.13.1 in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.0->peft==0.13.3.dev0) (1.13.1)\n",
            "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from sympy==1.13.1->torch>=1.13.0->peft==0.13.3.dev0) (1.3.0)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from transformers->peft==0.13.3.dev0) (2024.9.11)\n",
            "Requirement already satisfied: tokenizers<0.21,>=0.20 in /usr/local/lib/python3.10/dist-packages (from transformers->peft==0.13.3.dev0) (0.20.3)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch>=1.13.0->peft==0.13.3.dev0) (3.0.2)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->huggingface_hub>=0.25.0->peft==0.13.3.dev0) (3.4.0)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->huggingface_hub>=0.25.0->peft==0.13.3.dev0) (3.10)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->huggingface_hub>=0.25.0->peft==0.13.3.dev0) (2.2.3)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->huggingface_hub>=0.25.0->peft==0.13.3.dev0) (2024.8.30)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Imports"
      ],
      "metadata": {
        "id": "5BerCvfkq_jp"
      },
      "id": "5BerCvfkq_jp"
    },
    {
      "cell_type": "code",
      "source": [
        "from typing import Any, Dict, List, Union\n",
        "\n",
        "import numpy as np\n",
        "import torch\n",
        "from datasets import load_dataset\n",
        "from torch.utils.data import Dataset\n",
        "from tqdm import tqdm\n",
        "from transformers import (\n",
        "    AutoModelForCausalLM,\n",
        "    AutoTokenizer,\n",
        "    DataCollatorForLanguageModeling,\n",
        "    Trainer,\n",
        "    TrainingArguments,\n",
        ")\n",
        "\n",
        "from peft import CPTConfig, get_peft_model\n",
        "\n",
        "\n",
        "MAX_INPUT_LENGTH = 1024\n",
        "MAX_ICL_SAMPLES = 10\n",
        "NUM_TRAINING_SAMPLES = 100\n",
        "model_id = 'bigscience/bloom-1b7'"
      ],
      "metadata": {
        "id": "Y0pETNFBl963"
      },
      "id": "Y0pETNFBl963",
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Data Preparation\n",
        "---"
      ],
      "metadata": {
        "id": "9hO_I3aDmCQu"
      },
      "id": "9hO_I3aDmCQu"
    },
    {
      "cell_type": "code",
      "source": [
        "# Initialize the tokenizer\n",
        "tokenizer = AutoTokenizer.from_pretrained(\n",
        "    model_id,               # The name or path of the pre-trained tokenizer (e.g., \"bert-base-uncased\").\n",
        "    cache_dir='.',          # Directory to cache the tokenizer files locally.\n",
        "    padding_side='right',   # Specifies that padding should be added to the right side of sequences.\n",
        "    trust_remote_code=True  # Allows loading tokenizer implementations from external sources.\n",
        ")"
      ],
      "metadata": {
        "id": "STK5N0LJrZmA",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "4c5c3dda-07ae-4f67-df29-4a2ff499e5ad"
      },
      "id": "STK5N0LJrZmA",
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/huggingface_hub/utils/_auth.py:94: UserWarning: \n",
            "The secret `HF_TOKEN` does not exist in your Colab secrets.\n",
            "To authenticate with the Hugging Face Hub, create a token in your settings tab (https://huggingface.co/settings/tokens), set it as secret in your Google Colab and restart your session.\n",
            "You will be able to reuse this secret in all of your notebooks.\n",
            "Please note that authentication is recommended but still optional to access public models or datasets.\n",
            "  warnings.warn(\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Load the SST-2 dataset from the GLUE benchmark\n",
        "dataset = load_dataset('glue', 'sst2')\n",
        "\n",
        "def add_string_labels(example):\n",
        "    \"\"\"\n",
        "    Converts numerical labels into human-readable string labels.\n",
        "\n",
        "    Args:\n",
        "        example (dict): A single example from the dataset with a numerical 'label'.\n",
        "\n",
        "    Returns:\n",
        "        dict: The example augmented with a 'label_text' field.\n",
        "    \"\"\"\n",
        "    # Map numerical label to string label\n",
        "    example['label_text'] = \"positive\" if example['label'] == 1 else \"negative\"\n",
        "    return example\n",
        "\n",
        "# Subset and process the training dataset\n",
        "context_dataset = dataset['train'].select(range(MAX_ICL_SAMPLES)).map(add_string_labels)\n",
        "train_dataset = dataset['train'].select(range(MAX_ICL_SAMPLES, NUM_TRAINING_SAMPLES + MAX_ICL_SAMPLES)).map(add_string_labels)"
      ],
      "metadata": {
        "id": "C3oq4lDDrcUf",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "referenced_widgets": [
            "72a5be4b77ec4d5994bcace9d462da84",
            "bed78529ff2c4d08befca97c50cb5efc",
            "cf7077acfce04aff8af0a2483dbf094c",
            "910462d70d944d00ba54958d77bee755",
            "a899818bdad0415b860eaac4afe31f30",
            "3d78a6c8923547cf8c75bc8c10125eda",
            "8083f95a673a423286ade63051de757d",
            "13fc203ab1b44c83b6cfcc1e171d26ad",
            "663a0196d2b547fd8a6890b8a86080c2",
            "72be01164e974d59b05bee716e9bc978",
            "4cedaf37e79e4ff1a10ffb96ec543e81"
          ],
          "height": 49
        },
        "outputId": "5ae1ff54-d726-4f07-e6d7-cd53145b5d6f"
      },
      "id": "C3oq4lDDrcUf",
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Map:   0%|          | 0/100 [00:00<?, ? examples/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "72a5be4b77ec4d5994bcace9d462da84"
            }
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Note:** This notebook uses small subsets of the dataset to ensure quick execution. For proper testing and evaluation, it is recommended to use the entire dataset by setting quick_review to False."
      ],
      "metadata": {
        "id": "ehlJCE80TnrC"
      },
      "id": "ehlJCE80TnrC"
    },
    {
      "cell_type": "code",
      "source": [
        "quick_review = True # set to False for a comprehensive evaluation\n",
        "num_of_test_examples = 100 if quick_review else len(dataset['validation'])\n",
        "# Subset and process the validation dataset\n",
        "test_dataset = dataset['validation'].select(range(num_of_test_examples)).map(add_string_labels)"
      ],
      "metadata": {
        "id": "yXoBN6EwTmNX"
      },
      "id": "yXoBN6EwTmNX",
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [],
      "metadata": {
        "id": "YpXEWzglTl74"
      },
      "id": "YpXEWzglTl74"
    },
    {
      "cell_type": "code",
      "source": [
        "class CPTDataset(Dataset):\n",
        "    def __init__(self, samples, tokenizer, template, max_length=MAX_INPUT_LENGTH):\n",
        "        \"\"\"\n",
        "        Initialize the CPTDataset with samples, a tokenizer, and a template.\n",
        "\n",
        "        Args:\n",
        "            samples (list): List of samples containing input sentences and labels.\n",
        "            tokenizer: Tokenizer instance for encoding text.\n",
        "            template (dict): Dictionary defining input/output templates and separators.\n",
        "            max_length (int): Maximum input length for truncation.\n",
        "        \"\"\"\n",
        "        self.template = template\n",
        "        self.tokenizer = tokenizer\n",
        "        self.max_length = max_length\n",
        "\n",
        "        # Storage for tokenized inputs and masks\n",
        "        self.attention_mask = []\n",
        "        self.input_ids = []\n",
        "        self.input_type_mask = []\n",
        "        self.inter_seperator_ids = self._get_input_ids(template['inter_seperator'])\n",
        "\n",
        "        # Tokenize each sample and prepare inputs\n",
        "        for sample_i in tqdm(samples):\n",
        "            input_text, label = sample_i['sentence'], sample_i['label_text']\n",
        "            input_ids, attention_mask, input_type_mask = self.preprocess_sentence(input_text, label)\n",
        "\n",
        "            self.input_ids.append(input_ids)\n",
        "            self.attention_mask.append(attention_mask)\n",
        "            self.input_type_mask.append(input_type_mask)\n",
        "\n",
        "\n",
        "    def _get_input_ids(self, text):\n",
        "        \"\"\"\n",
        "        Tokenize the given text into input IDs.\n",
        "\n",
        "        Args:\n",
        "            text (str): The text to tokenize.\n",
        "\n",
        "        Returns:\n",
        "            list: Tokenized input IDs.\n",
        "        \"\"\"\n",
        "        return self.tokenizer(text, add_special_tokens=False)[\"input_ids\"]\n",
        "\n",
        "\n",
        "    def preprocess_sentence(self, input_text, label):\n",
        "        \"\"\"\n",
        "        Preprocess a sentence and its corresponding label using templates.\n",
        "\n",
        "        Args:\n",
        "            input_text (str): The input sentence.\n",
        "            label (str): The label text (e.g., \"positive\", \"negative\").\n",
        "\n",
        "        Returns:\n",
        "            tuple: (input_ids, attention_mask, input_type_mask)\n",
        "        \"\"\"\n",
        "\n",
        "        # Split input template into parts\n",
        "        input_template_part_1_text, input_template_part_2_text = self.template['input'].split('{}')\n",
        "        input_template_tokenized_part1 = self._get_input_ids(input_template_part_1_text)\n",
        "        input_tokenized = self._get_input_ids(input_text)\n",
        "        input_template_tokenized_part2 = self._get_input_ids(input_template_part_2_text)\n",
        "\n",
        "        # Separator token\n",
        "        sep_tokenized = self._get_input_ids(self.template['intra_seperator'])\n",
        "\n",
        "        # Process the label using the template\n",
        "        label_template_part_1, label_template_part_2 = self.template['output'].split('{}')\n",
        "        label_template_part1_tokenized = self._get_input_ids(label_template_part_1)\n",
        "        label_tokenized = self._get_input_ids(label)\n",
        "        label_template_part2_tokenized = self._get_input_ids(label_template_part_2)\n",
        "\n",
        "        # End-of-sequence token\n",
        "        eos = [self.tokenizer.eos_token_id] if self.tokenizer.eos_token_id is not None else []\n",
        "\n",
        "        # Concatenate all tokenized parts\n",
        "        input_ids = input_template_tokenized_part1 + input_tokenized + input_template_tokenized_part2 + sep_tokenized + label_template_part1_tokenized + label_tokenized + label_template_part2_tokenized + eos\n",
        "\n",
        "        # Generate attention and type masks\n",
        "        attention_mask = [1] * len(input_ids)\n",
        "        input_type_mask = [1] * len(input_template_tokenized_part1) + [2] * len(input_tokenized) + [1] * len(\n",
        "            input_template_tokenized_part2) + [0] * len(sep_tokenized) + \\\n",
        "                          [3] * len(label_template_part1_tokenized) + [4] * len(label_tokenized) + [3] * len( \\\n",
        "            label_template_part2_tokenized) + [0] * len(eos)\n",
        "\n",
        "        # Ensure all masks and inputs are the same length\n",
        "        assert len(input_type_mask) == len(input_ids) == len(attention_mask)\n",
        "\n",
        "        return input_ids, attention_mask, input_type_mask\n",
        "\n",
        "\n",
        "    def __len__(self):\n",
        "        \"\"\"\n",
        "        Get the number of examples in the dataset.\n",
        "\n",
        "        Returns:\n",
        "            int: Number of examples.\n",
        "        \"\"\"\n",
        "        return len(self.input_ids)\n",
        "\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        \"\"\"\n",
        "        Get the tokenized representation for the given index.\n",
        "\n",
        "        Args:\n",
        "            idx (int): Index of the example.\n",
        "\n",
        "        Returns:\n",
        "            dict: Tokenized inputs with attention and type masks.\n",
        "        \"\"\"\n",
        "\n",
        "        return {\n",
        "            \"input_ids\": self.input_ids[idx],\n",
        "            \"attention_mask\": self.attention_mask[idx],\n",
        "            \"input_type_mask\": self.input_type_mask[idx]\n",
        "        }\n",
        "\n",
        "# Define templates for tokenization\n",
        "templates = {\n",
        "    'input': 'input: {}',     # Input template with placeholder\n",
        "    'intra_seperator': ' ',   # Separator between input and output\n",
        "    'output': 'output: {}',   # Output template with placeholder\n",
        "    'inter_seperator': '\\n'   # Separator between examples\n",
        "}\n",
        "\n",
        "# Initialize the dataset\n",
        "cpt_train_dataset = CPTDataset(train_dataset, tokenizer, templates)\n",
        "\n",
        "\n",
        "# - `templates`: Define how inputs and outputs should be formatted and separated.\n",
        "# - `CPTDataset`: Converts the raw dataset into tokenized input IDs and masks."
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "icJb6-Uqrf8p",
        "outputId": "a0a2ddbc-1d1f-4845-93c9-19cf34b46024"
      },
      "id": "icJb6-Uqrf8p",
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 100/100 [00:00<00:00, 874.85it/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 10/10 [00:00<00:00, 1133.50it/s]\n"
          ]
        }
      ],
      "source": [
        "# Initialize storage for context-level information\n",
        "context_ids = []                # Concatenated input IDs for all samples\n",
        "context_attention_mask = []     # Concatenated attention masks\n",
        "context_input_type_mask = []    # Concatenated input type masks\n",
        "first_type_mask = 0             # Initial offset for input type mask\n",
        "\n",
        "cpt_context_dataset = CPTDataset(context_dataset, tokenizer, templates)\n",
        "\n",
        "# Iterate through the CPT training dataset\n",
        "for i in range(len(context_dataset)):\n",
        "    # Add input IDs to the context\n",
        "    context_ids += cpt_context_dataset[i]['input_ids']\n",
        "\n",
        "    # Add attention mask to the context\n",
        "    context_attention_mask += cpt_context_dataset[i]['attention_mask']\n",
        "\n",
        "    # Adjust and add the input type mask to the context\n",
        "    context_input_type_mask += [\n",
        "        i + first_type_mask if i > 0 else 0 # Increment type indices dynamically\n",
        "        for i in cpt_context_dataset[i]['input_type_mask']\n",
        "        ]\n",
        "\n",
        "    # Increment the type mask offset after processing the sample\n",
        "    first_type_mask += 4"
      ],
      "metadata": {
        "ExecuteTime": {
          "end_time": "2024-10-22T09:24:58.894814Z",
          "start_time": "2024-10-22T09:24:58.893841Z"
        },
        "id": "aef03bbd5d86d3d8",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "1bb1343b-b5f8-4998-e34b-6a8ae8063381"
      },
      "id": "aef03bbd5d86d3d8",
      "execution_count": 7
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Model Training\n",
        "\n",
        "---"
      ],
      "metadata": {
        "collapsed": false,
        "id": "2c40f24774d83372"
      },
      "id": "2c40f24774d83372"
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Load model"
      ],
      "metadata": {
        "id": "p0jFTzkisMgN"
      },
      "id": "p0jFTzkisMgN"
    },
    {
      "cell_type": "code",
      "outputs": [],
      "source": [
        "# Load a pre-trained causal language model\n",
        "base_model = AutoModelForCausalLM.from_pretrained(\n",
        "    model_id,\n",
        "    cache_dir='.',\n",
        "    dtype=torch.float16,\n",
        "    device_map='auto'\n",
        ")\n",
        "\n",
        "# Initialize the CPT configuration\n",
        "config = CPTConfig(\n",
        "            cpt_token_ids=context_ids,\n",
        "            cpt_mask=context_attention_mask,\n",
        "            cpt_tokens_type_mask=context_input_type_mask,\n",
        "\n",
        "            opt_weighted_loss_type='decay',\n",
        "            opt_loss_decay_factor=0.95,         # we choose the exponential decay factor applied to the loss\n",
        "            opt_projection_epsilon=0.2,         # we choose the projection over the input tokens\n",
        "            opt_projection_format_epsilon=0.1,  # we choose the projection over input and output templates\n",
        "\n",
        "            tokenizer_name_or_path=model_id,\n",
        ")\n",
        "\n",
        "# Initialize the CPT model with PEFT\n",
        "model = get_peft_model(base_model, config)"
      ],
      "metadata": {
        "ExecuteTime": {
          "end_time": "2024-10-22T09:25:08.941945Z",
          "start_time": "2024-10-22T09:25:04.393323Z"
        },
        "id": "17ac445134919a39"
      },
      "id": "17ac445134919a39",
      "execution_count": 8
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Setting Collate Function"
      ],
      "metadata": {
        "collapsed": false,
        "id": "4e49660c50d98741"
      },
      "id": "4e49660c50d98741"
    },
    {
      "cell_type": "code",
      "outputs": [],
      "source": [
        "class CPTDataCollatorForLanguageModeling(DataCollatorForLanguageModeling):\n",
        "    def __init__(self, tokenizer, training=True, mlm=False):\n",
        "        \"\"\"\n",
        "        Custom collator for CPT-style language modeling.\n",
        "\n",
        "        Args:\n",
        "            tokenizer: The tokenizer to handle tokenization and special tokens.\n",
        "            training (bool): If True, operates in training mode; otherwise in evaluation mode.\n",
        "            mlm (bool): If True, enables masked language modeling.\n",
        "        \"\"\"\n",
        "\n",
        "        super().__init__(tokenizer, mlm=mlm) # Initialize the parent class\n",
        "        self.training = training\n",
        "\n",
        "        # Add a special padding token if not already defined\n",
        "        self.tokenizer.add_special_tokens({\"pad_token\": \"[PAD]\"})\n",
        "\n",
        "    def torch_call(self, examples: List[Union[List[int], Any, Dict[str, Any]]]) -> Dict[str, Any]:\n",
        "        \"\"\"\n",
        "        Process a batch of examples for language modeling.\n",
        "\n",
        "        Args:\n",
        "            examples (List): A batch of examples with tokenized inputs and optional sample masks.\n",
        "\n",
        "        Returns:\n",
        "            Dict: A dictionary containing padded and tensor-converted inputs, attention masks,\n",
        "                  input type masks, and optional sample masks and labels.\n",
        "        \"\"\"\n",
        "\n",
        "        # Initialize a list to collect sample masks if provided\n",
        "        list_sample_mask = []\n",
        "        for i in range(len(examples)):\n",
        "            if \"sample_mask\" in examples[i].keys():\n",
        "                list_sample_mask.append(examples[i].pop(\"sample_mask\"))\n",
        "\n",
        "        # Define a helper function for padding sequences to the maximum length\n",
        "        max_len = max(len(ex[\"input_ids\"]) for ex in examples)\n",
        "\n",
        "        # Define a helper function for padding sequences to the maximum length\n",
        "        def pad_sequence(sequence, max_len, pad_value=0):\n",
        "            return sequence + [pad_value] * (max_len - len(sequence))\n",
        "\n",
        "        # Pad and convert `input_ids`, `attention_mask`, and `input_type_mask` to tensors\n",
        "        input_ids = torch.tensor([pad_sequence(ex[\"input_ids\"], max_len) for ex in examples])\n",
        "        attention_mask = torch.tensor([pad_sequence(ex[\"attention_mask\"], max_len) for ex in examples])\n",
        "        input_type_mask = torch.tensor([pad_sequence(ex[\"input_type_mask\"], max_len) for ex in examples])\n",
        "\n",
        "        # Create the initial batch dictionary\n",
        "        batch = {\"input_ids\": input_ids, \"attention_mask\": attention_mask, \"input_type_mask\": input_type_mask}\n",
        "\n",
        "        # Create a tensor to store sample masks\n",
        "        tensor_sample_mask = batch[\"input_ids\"].clone().long()\n",
        "        tensor_sample_mask[:, :] = 0 # Initialize with zeros\n",
        "\n",
        "        # Populate the tensor with the provided sample masks\n",
        "        for i in range(len(list_sample_mask)):\n",
        "            tensor_sample_mask[i, : len(list_sample_mask[i])] = list_sample_mask[i]\n",
        "\n",
        "        # Copy `input_ids` to use as `labels`\n",
        "        batch[\"labels\"] = batch[\"input_ids\"].clone()\n",
        "\n",
        "        # If in evaluation mode, include the `sample_mask` in the batch\n",
        "        if not self.training:\n",
        "            batch[\"sample_mask\"] = tensor_sample_mask\n",
        "\n",
        "        return batch"
      ],
      "metadata": {
        "ExecuteTime": {
          "end_time": "2024-10-22T09:25:08.953199Z",
          "start_time": "2024-10-22T09:25:08.945689Z"
        },
        "id": "b0fac840f060e3aa"
      },
      "id": "b0fac840f060e3aa",
      "execution_count": 9
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Training"
      ],
      "metadata": {
        "collapsed": false,
        "id": "48f535d74e6602b"
      },
      "id": "48f535d74e6602b"
    },
    {
      "cell_type": "code",
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "    <div>\n",
              "      \n",
              "      <progress value='500' max='500' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
              "      [500/500 01:28, Epoch 5/5]\n",
              "    </div>\n",
              "    <table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              " <tr style=\"text-align: left;\">\n",
              "      <th>Step</th>\n",
              "      <th>Training Loss</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <td>100</td>\n",
              "      <td>0.400800</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>200</td>\n",
              "      <td>0.036000</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>300</td>\n",
              "      <td>0.026300</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>400</td>\n",
              "      <td>0.016100</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>500</td>\n",
              "      <td>0.011600</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table><p>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "TrainOutput(global_step=500, training_loss=0.09815525007247924, metrics={'train_runtime': 90.6767, 'train_samples_per_second': 5.514, 'train_steps_per_second': 5.514, 'total_flos': 79477977907200.0, 'train_loss': 0.09815525007247924, 'epoch': 5.0})"
            ]
          },
          "metadata": {},
          "execution_count": 10
        }
      ],
      "source": [
        "training_args = TrainingArguments(\n",
        "    output_dir='../.',\n",
        "    use_cpu=False,\n",
        "    auto_find_batch_size=False,\n",
        "    learning_rate=1e-4,\n",
        "    logging_steps=100,\n",
        "    per_device_train_batch_size=1,\n",
        "    save_total_limit=1,\n",
        "    remove_unused_columns=False,\n",
        "    num_train_epochs=5,\n",
        "    fp16=True,\n",
        "    save_strategy='no',\n",
        "    logging_dir=\"logs\",\n",
        "    report_to=\"none\"\n",
        ")\n",
        "\n",
        "trainer = Trainer(\n",
        "    model=model,\n",
        "    args=training_args,\n",
        "    train_dataset=cpt_train_dataset,  # Custom CPT training dataset.\n",
        "    data_collator=CPTDataCollatorForLanguageModeling(tokenizer, training=True, mlm=False)\n",
        ")\n",
        "\n",
        "trainer.train()"
      ],
      "metadata": {
        "ExecuteTime": {
          "end_time": "2024-10-22T09:25:27.599132Z",
          "start_time": "2024-10-22T09:25:13.906685Z"
        },
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 268
        },
        "id": "1a865c2ad2dc7218",
        "outputId": "c4bfd785-e354-4ee6-a87e-63c17bfd2605"
      },
      "id": "1a865c2ad2dc7218",
      "execution_count": 10
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Model Evaluation\n",
        "\n",
        "---"
      ],
      "metadata": {
        "collapsed": false,
        "id": "b799ea89a567590f"
      },
      "id": "b799ea89a567590f"
    },
    {
      "cell_type": "code",
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 100/100 [00:00<00:00, 1972.82it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Sentence: input: it 's a charming and often affecting journey .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: unflinchingly bleak and desperate  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: allows us to hope that nolan is poised to embark a major career as a commercial yet inventive filmmaker .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: the acting , costumes , music , cinematography and sound are all astounding given the production 's austere locales .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: it 's slow -- very , very slow .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: although laced with humor and a few fanciful touches , the film is a refreshingly serious look at young women .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: a sometimes tedious film .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: or doing last year 's taxes with your ex-wife .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: you do n't have to know about music to appreciate the film 's easygoing blend of comedy and romance .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: in exactly 89 minutes , most of which passed as slowly as if i 'd been sitting naked on an igloo , formula 51 sank from quirky to jerky to utter turkey .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: the mesmerizing performances of the leads keep the film grounded and keep the audience riveted .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: it takes a strange kind of laziness to waste the talents of robert forster , anne meara , eugene levy , and reginald veljohnson all in the same movie .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: ... the film suffers from a lack of humor ( something needed to balance out the violence ) ...  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: we root for ( clara and paul ) , even like them , though perhaps it 's an emotion closer to pity .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: even horror fans will most likely not find what they 're seeking with trouble every day ; the movie lacks both thrills and humor .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: a gorgeous , high-spirited musical from india that exquisitely blends music , dance , song , and high drama .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: the emotions are raw and will strike a nerve with anyone who 's ever had family trauma .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: audrey tatou has a knack for picking roles that magnify her outrageous charm , and in this literate french comedy , she 's as morning-glory exuberant as she was in amélie .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: ... the movie is just a plain old monster .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: in its best moments , resembles a bad high school production of grease , without benefit of song .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: pumpkin takes an admirable look at the hypocrisy of political correctness , but it does so with such an uneven tone that you never know when humor ends and tragedy begins .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: the iditarod lasts for days - this just felt like it did .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: holden caulfield did it better .  output: negative</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is negative\n",
            "Sentence: input: a delectable and intriguing thriller filled with surprises , read my lips is an original .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: seldom has a movie so closely matched the spirit of a man and his work .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: nicks , seemingly uncertain what 's going to make people laugh , runs the gamut from stale parody to raunchy sex gags to formula romantic comedy .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: the action switches between past and present , but the material link is too tenuous to anchor the emotional connections that purport to span a 125-year divide .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: it 's an offbeat treat that pokes fun at the democratic exercise while also examining its significance for those who take part .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: it 's a cookie-cutter movie , a cut-and-paste job .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: i had to look away - this was god awful .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: thanks to scott 's charismatic roger and eisenberg 's sweet nephew , roger dodger is one of the most compelling variations on in the company of men .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: ... designed to provide a mix of smiles and tears , `` crossroads '' instead provokes a handful of unintentional howlers and numerous yawns .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: a gorgeous , witty , seductive movie .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: if the movie succeeds in instilling a wary sense of ` there but for the grace of god , ' it is far too self-conscious to draw you deeply into its world .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: it does n't believe in itself , it has no sense of humor ... it 's just plain bored .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: a sequence of ridiculous shoot - 'em - up scenes .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: the weight of the piece , the unerring professionalism of the chilly production , and the fascination embedded in the lurid topic prove recommendation enough .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: ( w ) hile long on amiable monkeys and worthy environmentalism , jane goodall 's wild chimpanzees is short on the thrills the oversize medium demands .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: as surreal as a dream and as detailed as a photograph , as visually dexterous as it is at times imaginatively overwhelming .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: escaping the studio , piccoli is warmly affecting and so is this adroitly minimalist movie .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: there 's ... tremendous energy from the cast , a sense of playfulness and excitement that seems appropriate .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: this illuminating documentary transcends our preconceived vision of the holy land and its inhabitants , revealing the human complexities beneath .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: the subtle strength of `` elling '' is that it never loses touch with the reality of the grim situation .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: holm ... embodies the character with an effortlessly regal charisma .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: the title not only describes its main characters , but the lazy people behind the camera as well .  output: negative</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is negative\n",
            "Sentence: input: it offers little beyond the momentary joys of pretty and weightless intellectual entertainment .  output: negative</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is negative\n",
            "Sentence: input: a synthesis of cliches and absurdities that seems positively decadent in its cinematic flash and emptiness .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: a subtle and well-crafted ( for the most part ) chiller .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: has a lot of the virtues of eastwood at his best .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: it 's hampered by a lifetime-channel kind of plot and a lead actress who is out of her depth .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: it feels like an after-school special gussied up with some fancy special effects , and watching its rote plot points connect is about as exciting as gazing at an egg timer for 93 minutes .  output: negative</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is negative\n",
            "Sentence: input: for the most part , director anne-sophie birot 's first feature is a sensitive , extraordinarily well-acted drama .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: mr. tsai is a very original artist in his medium , and what time is it there ?  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: sade is an engaging look at the controversial eponymous and fiercely atheistic hero .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: so devoid of any kind of intelligible story that it makes films like xxx and collateral damage seem like thoughtful treatises  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: a tender , heartfelt family drama .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: ... a hollow joke told by a cinematic gymnast having too much fun embellishing the misanthropic tale to actually engage it .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: the cold turkey would 've been a far better title .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: manages to be both repulsively sadistic and mundane .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: it 's just disappointingly superficial -- a movie that has all the elements necessary to be a fascinating , involving character study , but never does more than scratch the surface .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: this is a story of two misfits who do n't stand a chance alone , but together they are magnificent .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: schaeffer has to find some hook on which to hang his persistently useless movies , and it might as well be the resuscitation of the middle-aged character .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: the primitive force of this film seems to bubble up from the vast collective memory of the combatants .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: on this tricky topic , tadpole is very much a step in the right direction , with its blend of frankness , civility and compassion .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: the script kicks in , and mr. hartley 's distended pace and foot-dragging rhythms follow .  output: negative</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is negative\n",
            "Sentence: input: you wonder why enough was n't just a music video rather than a full-length movie .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: if you 're hard up for raunchy college humor , this is your ticket right here .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: a fast , funny , highly enjoyable movie .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: good old-fashioned slash-and-hack is back !  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: this one is definitely one to skip , even for horror movie fanatics .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: for all its impressive craftsmanship , and despite an overbearing series of third-act crescendos , lily chou-chou never really builds up a head of emotional steam .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: exquisitely nuanced in mood tics and dialogue , this chamber drama is superbly acted by the deeply appealing veteran bouquet and the chilling but quite human berling .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: uses high comedy to evoke surprising poignance .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: one of creepiest , scariest movies to come along in a long , long time , easily rivaling blair witch or the others .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: a string of rehashed sight gags based in insipid vulgarity .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: among the year 's most intriguing explorations of alientation .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: the movie fails to live up to the sum of its parts .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: the son 's room is a triumph of gentility that earns its moments of pathos .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: there is nothing outstanding about this film , but it is good enough and will likely be appreciated most by sailors and folks who know their way around a submarine .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: this is a train wreck of an action film -- a stupefying attempt by the filmmakers to force-feed james bond into the mindless xxx mold and throw 40 years of cinematic history down the toilet in favor of bright flashes and loud bangs .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: the draw ( for `` big bad love '' ) is a solid performance by arliss howard .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: green might want to hang onto that ski mask , as robbery may be the only way to pay for his next project .  output: negative</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is negative\n",
            "Sentence: input: it 's one pussy-ass world when even killer-thrillers revolve around group therapy sessions .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: though it 's become almost redundant to say so , major kudos go to leigh for actually casting people who look working-class .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: the band 's courage in the face of official repression is inspiring , especially for aging hippies ( this one included ) .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: the movie achieves as great an impact by keeping these thoughts hidden as ... ( quills ) did by showing them .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: the film flat lines when it should peak and is more missed opportunity and trifle than dark , decadent truffle .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: jaglom ... put ( s ) the audience in the privileged position of eavesdropping on his characters  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: fresnadillo 's dark and jolting images have a way of plying into your subconscious like the nightmare you had a week ago that wo n't go away .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: we know the plot 's a little crazy , but it held my interest from start to finish .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: it 's a scattershot affair , but when it hits its mark it 's brilliant .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: hardly a masterpiece , but it introduces viewers to a good charitable enterprise and some interesting real people .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: you wo n't like roger , but you will quickly recognize him .  output: negative</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is negative\n",
            "Sentence: input: if steven soderbergh 's ` solaris ' is a failure it is a glorious failure .  output: positive</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is positive\n",
            "Sentence: input: byler reveals his characters in a way that intrigues and even fascinates us , and he never reduces the situation to simple melodrama .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: this riveting world war ii moral suspense story deals with the shadow side of american culture : racial prejudice in its ugly and diverse forms .  output: negative</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is negative\n",
            "Sentence: input: it 's difficult to imagine the process that produced such a script , but here 's guessing that spray cheese and underarm noises played a crucial role .  output: negative</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is negative\n",
            "Sentence: input: no sophomore slump for director sam mendes , who segues from oscar winner to oscar-winning potential with a smooth sleight of hand .  output: positive</s> \n",
            " \t The prediction is: positive\n",
            " \t The GT is positive\n",
            "Sentence: input: on the whole , the movie lacks wit , feeling and believability to compensate for its incessant coarseness and banality .  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "Sentence: input: why make a documentary about these marginal historical figures ?  output: negative</s> \n",
            " \t The prediction is: negative\n",
            " \t The GT is negative\n",
            "The model Acc is 90.0%\n"
          ]
        }
      ],
      "source": [
        "model.eval()\n",
        "\n",
        "# Select relevant columns from the test dataset\n",
        "test_dataset = test_dataset.select_columns(['sentence', 'label_text'])\n",
        "\n",
        "# Convert the test dataset to a CPT-compatible format\n",
        "cpt_test_dataset = CPTDataset(test_dataset, tokenizer, templates)\n",
        "\n",
        "# Get the device where the model is loaded (CPU, GPU or XPU)\n",
        "device = model.device\n",
        "list_bool_predictions = []\n",
        "\n",
        "for i in range(len(test_dataset)):\n",
        "    input_ids, input_type_mask = cpt_test_dataset[i]['input_ids'], cpt_test_dataset[i]['input_type_mask']\n",
        "\n",
        "    # Pass the inputs through the model\n",
        "    outputs = model(\n",
        "        input_ids=torch.Tensor(input_ids).long().to(device=device).view(1, -1),\n",
        "        labels=torch.Tensor(input_ids).long().to(device=device).view(1, -1),\n",
        "        input_type_mask=torch.Tensor(input_type_mask).long().to(device=device).view(1, -1)\n",
        "    )\n",
        "\n",
        "    # Shift logits to exclude the last token and match the labels\n",
        "    shifted_logits = outputs.logits[..., :-1, :].contiguous().to(model.dtype)[0, -len(input_ids) + 1:]\n",
        "    shift_labels = torch.Tensor(input_ids).long().to(device=device).view(1, -1)[0, 1:].contiguous().to(device)\n",
        "    shifted_input_type_mask = torch.Tensor(input_type_mask).long().to(device=device).view(1, -1)[..., 1:].contiguous().to(device)\n",
        "\n",
        "    # Create a mask for the type `4` tokens (label tokens)\n",
        "    mask = torch.Tensor(shifted_input_type_mask).long().to(device=device).view(-1,) == 4\n",
        "\n",
        "    # Extract logits and labels corresponding to the mask\n",
        "    logit = shifted_logits[mask]\n",
        "    label = shift_labels[mask]\n",
        "\n",
        "    # All possible label tokens for `negative` and `positive`\n",
        "    all_labels = torch.Tensor([tokenizer(i, add_special_tokens=False)[\"input_ids\"] for i in ['negative', 'positive']]).long().to(device).view(-1,)\n",
        "\n",
        "    # Compare logits with label tokens and infer prediction\n",
        "    prediction = logit[0, torch.Tensor([tokenizer(i, add_special_tokens=False)[\"input_ids\"] for i in ['negative', 'positive']]).long().to(device).view(-1,)].argmax()\n",
        "    prediction_text = 'negative' if prediction == 0 else 'positive'\n",
        "    print(f\"Sentence: {tokenizer.decode(input_ids)} \\n \\t The prediction is: {prediction_text}\\n \\t The GT is {tokenizer.decode(label)}\")\n",
        "    list_bool_predictions.append(prediction_text == tokenizer.decode(label))\n",
        "\n",
        "print(f'The model Acc is {100 * np.mean(list_bool_predictions)}%')"
      ],
      "metadata": {
        "ExecuteTime": {
          "end_time": "2024-10-22T09:25:28.252009Z",
          "start_time": "2024-10-22T09:25:27.598326Z"
        },
        "id": "48e7d976e6e01212",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "40dd1226-fa31-4e77-dc7e-e06a3600304e"
      },
      "id": "48e7d976e6e01212",
      "execution_count": 11
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 2
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython2",
      "version": "2.7.6"
    },
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "72a5be4b77ec4d5994bcace9d462da84": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_bed78529ff2c4d08befca97c50cb5efc",
              "IPY_MODEL_cf7077acfce04aff8af0a2483dbf094c",
              "IPY_MODEL_910462d70d944d00ba54958d77bee755"
            ],
            "layout": "IPY_MODEL_a899818bdad0415b860eaac4afe31f30"
          }
        },
        "bed78529ff2c4d08befca97c50cb5efc": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_3d78a6c8923547cf8c75bc8c10125eda",
            "placeholder": "​",
            "style": "IPY_MODEL_8083f95a673a423286ade63051de757d",
            "value": "Map: 100%"
          }
        },
        "cf7077acfce04aff8af0a2483dbf094c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_13fc203ab1b44c83b6cfcc1e171d26ad",
            "max": 100,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_663a0196d2b547fd8a6890b8a86080c2",
            "value": 100
          }
        },
        "910462d70d944d00ba54958d77bee755": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_72be01164e974d59b05bee716e9bc978",
            "placeholder": "​",
            "style": "IPY_MODEL_4cedaf37e79e4ff1a10ffb96ec543e81",
            "value": " 100/100 [00:00&lt;00:00, 1327.06 examples/s]"
          }
        },
        "a899818bdad0415b860eaac4afe31f30": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "3d78a6c8923547cf8c75bc8c10125eda": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "8083f95a673a423286ade63051de757d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "13fc203ab1b44c83b6cfcc1e171d26ad": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "663a0196d2b547fd8a6890b8a86080c2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "72be01164e974d59b05bee716e9bc978": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "4cedaf37e79e4ff1a10ffb96ec543e81": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}
