{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Setup: Importing Libraries\n",
        "In this cell, we import the necessary libraries for model training and tokenization. \n",
        "We use `torch` for handling the neural network, `transformers` for BERT tokenization and classification, and `sklearn` for splitting the dataset."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/NillionNetwork/aivm/blob/main/examples/2a-fine-tuning-bert-tiny.ipynb)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "!pip install datasets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6oSxXkJN2eUV",
        "outputId": "1ef90c68-10a9-4e08-9311-5648a005c6c0"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "from transformers import AutoTokenizer, AutoModelForSequenceClassification\n",
        "from transformers import AdamW, get_linear_schedule_with_warmup\n",
        "from sklearn.model_selection import train_test_split\n",
        "import numpy as np\n",
        "import time\n",
        "if torch.cuda.is_available():\n",
        "    print(f\"GPU: {torch.cuda.get_device_name(0)} is available.\")\n",
        "else:\n",
        "    print(\"No GPU available. Training will run on CPU.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Installing Required Packages\n",
        "Here we install the additional necessary packages, such as `datasets` for handling dataset loading, and `onnx` for exporting the trained PyTorch model into ONNX format for deployment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "f9vZeMA-P6j9",
        "outputId": "0416013d-64ba-489d-bf74-ef2d40a99798"
      },
      "outputs": [],
      "source": [
        "!pip install datasets onnx"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Custom Dataset Class\n",
        "This cell defines a `CustomDataset` class to handle the text and label processing.\n",
        "The dataset will tokenize each text, ensuring it has a fixed maximum length, and return the corresponding tokenized input IDs, attention mask, and labels."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "NjvBucTO2oft"
      },
      "outputs": [],
      "source": [
        "class CustomDataset(Dataset):\n",
        "    def __init__(self, texts, labels, tokenizer, max_length=128):\n",
        "        self.texts = texts\n",
        "        self.labels = labels\n",
        "        self.tokenizer = tokenizer\n",
        "        self.max_length = max_length\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.texts)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        text = str(self.texts[idx])\n",
        "        label = self.labels[idx]\n",
        "\n",
        "        encoding = self.tokenizer(\n",
        "            text,\n",
        "            add_special_tokens=True,\n",
        "            max_length=self.max_length,\n",
        "            padding='max_length',\n",
        "            truncation=True,\n",
        "            return_tensors='pt'\n",
        "        )\n",
        "\n",
        "        return {\n",
        "            'input_ids': encoding['input_ids'].flatten(),\n",
        "            'attention_mask': encoding['attention_mask'].flatten(),\n",
        "            'labels': torch.tensor(label, dtype=torch.long)\n",
        "        }\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Freezing BERT Layers\n",
        "The following function freezes the first `num_layers_to_freeze` layers of the BERT model during training. This is useful when fine-tuning models, as it helps to focus on training the classifier head without modifying the underlying pretrained BERT model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "vmfx9ZycPCiv"
      },
      "outputs": [],
      "source": [
        "def freeze_bert_layers(model, num_layers_to_freeze):\n",
        "    \"\"\"\n",
        "    Freezes the specified number of layers from the bottom of the BERT model.\n",
        "    Args:\n",
        "        model: The BERT model\n",
        "        num_layers_to_freeze: Number of layers to freeze (counting from bottom)\n",
        "    \"\"\"\n",
        "    # Freeze embeddings\n",
        "    for param in model.bert.embeddings.parameters():\n",
        "        param.requires_grad = False\n",
        "\n",
        "    # Freeze the specified number of encoder layers\n",
        "    for layer in model.bert.encoder.layer[:num_layers_to_freeze]:\n",
        "        for param in layer.parameters():\n",
        "            param.requires_grad = False\n",
        "\n",
        "    # Print trainable parameters info\n",
        "    total_params = sum(p.numel() for p in model.parameters())\n",
        "    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
        "    print(f\"Total parameters: {total_params:,}\")\n",
        "    print(f\"Trainable parameters: {trainable_params:,}\")\n",
        "    print(f\"Percentage of trainable parameters: {100 * trainable_params / total_params:.2f}%\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Training Loop\n",
        "This function handles training of the model for a specified number of epochs.\n",
        "It uses the AdamW optimizer and a linear learning rate scheduler to adjust the learning rate over time. The model is trained on batches of data and validated on a separate validation dataset after each epoch."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "pJjvSWdmPClb"
      },
      "outputs": [],
      "source": [
        "def train_model(model, train_loader, val_loader, device, num_epochs=3):\n",
        "    # Initialize optimizer only with parameters that require gradients\n",
        "    optimizer = AdamW(\n",
        "        [p for p in model.parameters() if p.requires_grad],\n",
        "        lr=2e-5\n",
        "    )\n",
        "\n",
        "    total_steps = len(train_loader) * num_epochs\n",
        "    scheduler = get_linear_schedule_with_warmup(\n",
        "        optimizer,\n",
        "        num_warmup_steps=0,\n",
        "        num_training_steps=total_steps\n",
        "    )\n",
        "\n",
        "    for epoch in range(num_epochs):\n",
        "        start_time = time.time()\n",
        "        model.train()\n",
        "        total_train_loss = 0\n",
        "        right_predictions = 0\n",
        "        for batch in train_loader:\n",
        "            input_ids = batch['input_ids'].to(device)\n",
        "            attention_mask = batch['attention_mask'].to(device)\n",
        "            labels = batch['labels'].to(device)\n",
        "\n",
        "            model.zero_grad()\n",
        "            outputs = model(\n",
        "                input_ids=input_ids,\n",
        "                attention_mask=attention_mask,\n",
        "                labels=labels\n",
        "            )\n",
        "\n",
        "            # Calculate accuracy\n",
        "            preds = torch.argmax(outputs.logits, dim=1)\n",
        "            right_predictions += torch.sum(preds == labels).item()\n",
        "\n",
        "            loss = outputs.loss\n",
        "            total_train_loss += loss.item()\n",
        "\n",
        "            loss.backward()\n",
        "            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
        "            optimizer.step()\n",
        "            scheduler.step()\n",
        "\n",
        "        # Validation\n",
        "        model.eval()\n",
        "        total_val_loss = 0\n",
        "        predictions = []\n",
        "        true_labels = []\n",
        "\n",
        "        with torch.no_grad():\n",
        "            for batch in val_loader:\n",
        "                input_ids = batch['input_ids'].to(device)\n",
        "                attention_mask = batch['attention_mask'].to(device)\n",
        "                labels = batch['labels'].to(device)\n",
        "\n",
        "                outputs = model(\n",
        "                    input_ids=input_ids,\n",
        "                    attention_mask=attention_mask,\n",
        "                    labels=labels\n",
        "                )\n",
        "\n",
        "                loss = outputs.loss\n",
        "                total_val_loss += loss.item()\n",
        "\n",
        "                preds = torch.argmax(outputs.logits, dim=1)\n",
        "                predictions.extend(preds.cpu().numpy())\n",
        "                true_labels.extend(labels.cpu().numpy())\n",
        "\n",
        "        avg_train_loss = total_train_loss / len(train_loader)\n",
        "        avg_val_loss = total_val_loss / len(val_loader)\n",
        "        accuracy = np.mean(np.array(predictions) == np.array(true_labels))\n",
        "        end_time = time.time()\n",
        "\n",
        "        print(f'Epoch {epoch + 1}:')\n",
        "        print(f'Average training loss: {avg_train_loss:.4f}')\n",
        "        print(f'Average validation loss: {avg_val_loss:.4f}')\n",
        "        print(f'Right predictions: {right_predictions} out of {len(train_loader) * 32}')\n",
        "        print(f'Validation Accuracy: {accuracy:.4f}')\n",
        "        print(f'Time taken for epoch: {end_time - start_time:.2f} seconds')\n",
        "        print('-' * 60)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Converting PyTorch Model to ONNX\n",
        "The next function converts the trained PyTorch model into ONNX format, a format optimized for deployment and inference."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "fQSA6XqaQsPC"
      },
      "outputs": [],
      "source": [
        "import tempfile, onnx\n",
        "\n",
        "\n",
        "def convert_pytorch_to_onnx_with_tokenizer(model, tokenizer, max_length=128, onnx_file_path=None):\n",
        "    \"\"\"\n",
        "    Converts a PyTorch model to ONNX format, using tokenizer output as input.\n",
        "\n",
        "    Args:\n",
        "    model (torch.nn.Module): The PyTorch model to be converted.\n",
        "    tokenizer: The tokenizer used to preprocess the input.\n",
        "    onnx_file_path (str): The file path where the ONNX model will be saved.\n",
        "    max_length (int): Maximum sequence length for the tokenizer.\n",
        "\n",
        "    Returns:\n",
        "    None\n",
        "    \"\"\"\n",
        "    model.eval()\n",
        "\n",
        "    # Prepare dummy input using the tokenizer\n",
        "    dummy_input = \"This is a sample input text for ONNX conversion.\"\n",
        "    inputs = tokenizer(\n",
        "        dummy_input,\n",
        "        max_length=max_length,\n",
        "        padding=\"max_length\",\n",
        "        truncation=True,\n",
        "        return_tensors=\"pt\",\n",
        "    )\n",
        "\n",
        "    # # Get the input names\n",
        "    input_names = list(inputs.keys())\n",
        "    input_names = [\"input_ids\", \"attention_mask\"]\n",
        "    print(f\"Input names: {input_names}\")\n",
        "\n",
        "    # # Create dummy inputs for ONNX export\n",
        "    # dummy_inputs = tuple(encoded_input[name] for name in input_names)\n",
        "    if onnx_file_path is None:\n",
        "      onnx_file_path = tempfile.mktemp(suffix=\".onnx\")\n",
        "    dynamic_axes = {name: {0: \"batch_size\"} for name in input_names}\n",
        "    dynamic_axes.update({f\"logits\": {0: \"batch_size\"}})\n",
        "    print(f\"dynamic_axes: {dynamic_axes}\")\n",
        "    # Export the model\n",
        "    torch.onnx.export(\n",
        "        model,  # model being run\n",
        "        tuple(inputs[k] for k in input_names),  # model inputs\n",
        "        onnx_file_path,  # where to save the model\n",
        "        export_params=True,  # store the trained parameter weights inside the model file\n",
        "        opset_version=20,  # the ONNX version to export the model to\n",
        "        do_constant_folding=True,  # whether to execute constant folding for optimization\n",
        "        input_names=input_names,  # the model's input names\n",
        "        output_names=[\"logits\"],  # the model's output names\n",
        "        dynamic_axes=dynamic_axes,\n",
        "    )  # variable length axes\n",
        "\n",
        "    print(f\"Model exported to {onnx_file_path}\")\n",
        "\n",
        "    # Verify the exported model\n",
        "    onnx_model = onnx.load(onnx_file_path)\n",
        "    onnx.checker.check_model(onnx_model)\n",
        "    print(\"ONNX model is valid.\")\n",
        "    return onnx_file_path, input_names\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Model and Tokenizer Initialization\n",
        "In this step, we initialize the `AutoTokenizer` and `AutoModelForSequenceClassification` using a pre-trained `bert-tiny` model. This model will be fine-tuned for binary sentiment classification."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XWicUYhJV0O4",
        "outputId": "d61efad3-b8c1-4955-ead1-ea37105392a2"
      },
      "outputs": [],
      "source": [
        "# Initialize tokenizer and model\n",
        "model_name = \"prajjwal1/bert-tiny\"\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
        "model = AutoModelForSequenceClassification.from_pretrained(\n",
        "    model_name,\n",
        "    num_labels=2  # Binary classification for sentiment\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Model Training\n",
        "\n",
        "This code you shared integrates several key functionalities to fine-tune a BERT model for sentiment analysis on the IMDB dataset, converts it to ONNX format, and evaluates it on example sentences. Here’s a refined and more organized version:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Pe2n6nEcPCnq",
        "outputId": "4e21734b-770f-49b9-df40-fc7e9cf7bb21"
      },
      "outputs": [],
      "source": [
        "def main(model, tokenizer):\n",
        "    # Load IMDB dataset\n",
        "    dataset = load_dataset(\"stanfordnlp/imdb\")\n",
        "\n",
        "    # Prepare train and validation datasets\n",
        "    train_texts = dataset['train']['text']\n",
        "    train_labels = dataset['train']['label']\n",
        "\n",
        "    # Split training data to create a validation set\n",
        "    train_texts, val_texts, train_labels, val_labels = train_test_split(\n",
        "        train_texts, train_labels, test_size=0.2, random_state=42, stratify=train_labels\n",
        "    )\n",
        "\n",
        "    # Create datasets\n",
        "    train_dataset = CustomDataset(train_texts, train_labels, tokenizer)\n",
        "    val_dataset = CustomDataset(val_texts, val_labels, tokenizer)\n",
        "\n",
        "    # Create dataloaders\n",
        "    train_loader = DataLoader(\n",
        "        train_dataset, batch_size=32, shuffle=True, num_workers=2\n",
        "    )\n",
        "    val_loader = DataLoader(\n",
        "        val_dataset, batch_size=32, num_workers=2\n",
        "    )\n",
        "\n",
        "    # Set device and move data to GPU if available\n",
        "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "    model.to(device)\n",
        "\n",
        "    # Train the model\n",
        "    train_model(model, train_loader, val_loader, device, num_epochs=3)\n",
        "\n",
        "    # Make model tensors contiguous and move to CPU before saving\n",
        "    model = model.cpu()\n",
        "\n",
        "    # Save the fine-tuned model as an ONNX file\n",
        "    onnx_file_path, input_names = convert_pytorch_to_onnx_with_tokenizer(\n",
        "        model, tokenizer, max_length=128, onnx_file_path=\"./my_new_bert_model.onnx\"\n",
        "    )\n",
        "    print(f\"ONNX file path: {onnx_file_path}\")\n",
        "    print(f\"Input names: {input_names}\")\n",
        "\n",
        "    # Test the model on a few examples\n",
        "    model.eval()\n",
        "    test_texts = [\n",
        "        \"This movie was absolutely fantastic! I loved every minute of it.\",\n",
        "        \"Terrible waste of time. The plot made no sense and the acting was awful.\"\n",
        "    ]\n",
        "\n",
        "    with torch.no_grad():\n",
        "        inputs = tokenizer(\n",
        "            test_texts,\n",
        "            padding=True,\n",
        "            truncation=True,\n",
        "            max_length=128,\n",
        "            return_tensors=\"pt\"\n",
        "        ).to(device)\n",
        "        model.to(device)\n",
        "        outputs = model(**inputs)\n",
        "        predictions = torch.nn.functional.softmax(outputs.logits, dim=1)\n",
        "\n",
        "        for text, pred in zip(test_texts, predictions):\n",
        "            sentiment = \"Positive\" if pred[1] > pred[0] else \"Negative\"\n",
        "            confidence = max(pred[0], pred[1]).item()\n",
        "            print(f\"\\nText: {text}\")\n",
        "            print(f\"Sentiment: {sentiment} (confidence: {confidence:.2%})\")\n",
        "\n",
        "# Call the main function\n",
        "main(model, tokenizer)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.12.6"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
