Adapter_Layers_for_Parameter_Efficient_Fine_Tuning.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "authorship_tag": "ABX9TyNM7t7xmroDTNiE97UIbJDw",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/OneFineStarstuff/OneFineStardust/blob/main/Adapter_Layers_for_Parameter_Efficient_Fine_Tuning.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import random\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from transformers import BertModel, BertTokenizer\n",
        "from torch.optim import AdamW\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "from torch.utils.tensorboard import SummaryWriter\n",
        "from sklearn.metrics import precision_score, recall_score, f1_score\n",
        "from nltk.corpus import wordnet\n",
        "from torch.optim.lr_scheduler import LambdaLR\n",
        "\n",
        "# Device configuration\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "# Define a custom dataset\n",
        "class TextDataset(Dataset):\n",
        "    def __init__(self, data, tokenizer, max_length=128, for_classification=False):\n",
        "        self.data = data\n",
        "        self.tokenizer = tokenizer\n",
        "        self.max_length = max_length\n",
        "        self.for_classification = for_classification\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.data)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        text = self.data[idx][\"text\"]\n",
        "        encoding = self.tokenizer(text, padding='max_length', truncation=True, max_length=self.max_length, return_tensors='pt')\n",
        "        input_ids = encoding[\"input_ids\"].squeeze()\n",
        "        attention_mask = encoding[\"attention_mask\"].squeeze()\n",
        "\n",
        "        if self.for_classification:\n",
        "            label = self.data[idx][\"label\"]\n",
        "            return input_ids, attention_mask, label\n",
        "        else:\n",
        "            return input_ids, attention_mask\n",
        "\n",
        "# Define FoundationModel class\n",
        "class FoundationModel(nn.Module):\n",
        "    def __init__(self, model_name=\"bert-base-uncased\", dropout_rate=0.1):\n",
        "        super(FoundationModel, self).__init__()\n",
        "        self.model = BertModel.from_pretrained(model_name)\n",
        "        self.tokenizer = BertTokenizer.from_pretrained(model_name)\n",
        "        self.dropout = nn.Dropout(dropout_rate)\n",
        "\n",
        "    def forward(self, input_ids, attention_mask):\n",
        "        outputs = self.model(input_ids, attention_mask=attention_mask)\n",
        "        return self.dropout(outputs.last_hidden_state)\n",
        "\n",
        "    def encode_text(self, texts, max_length=128):\n",
        "        encoding = self.tokenizer(texts, padding=True, truncation=True,\n",
        "                                  max_length=max_length, return_tensors=\"pt\")\n",
        "        return encoding[\"input_ids\"], encoding[\"attention_mask\"]\n",
        "\n",
        "# Define Adapter module\n",
        "class Adapter(nn.Module):\n",
        "    def __init__(self, input_dim, adapter_dim=64):\n",
        "        super(Adapter, self).__init__()\n",
        "        self.down_proj = nn.Linear(input_dim, adapter_dim)\n",
        "        self.up_proj = nn.Linear(adapter_dim, input_dim)\n",
        "        self.activation = nn.ReLU()\n",
        "        self.dropout = nn.Dropout(0.1)  # Add dropout to the adapter\n",
        "\n",
        "    def forward(self, x):\n",
        "        return x + self.dropout(self.up_proj(self.activation(self.down_proj(x))))\n",
        "\n",
        "# Integrate Adapter into Foundation Model\n",
        "class AdapterFoundationModel(FoundationModel):\n",
        "    def __init__(self, model_name=\"bert-base-uncased\", adapter_dim=64, dropout_rate=0.1):\n",
        "        super().__init__(model_name, dropout_rate)\n",
        "        for layer in self.model.encoder.layer:\n",
        "            layer.adapter = Adapter(self.model.config.hidden_size, adapter_dim)\n",
        "\n",
        "    def forward(self, input_ids, attention_mask):\n",
        "        outputs = self.model(input_ids, attention_mask=attention_mask)\n",
        "        for layer in self.model.encoder.layer:\n",
        "            layer.output = layer.adapter(layer.output)\n",
        "        return outputs\n",
        "\n",
        "# Define the MultiTaskAdapterFoundationModel class for multitask learning with adapters\n",
        "class MultiTaskAdapterFoundationModel(AdapterFoundationModel):\n",
        "    def __init__(self, model_name=\"bert-base-uncased\", tasks=None, adapter_dim=64, dropout_rate=0.1):\n",
        "        super().__init__(model_name, adapter_dim, dropout_rate)\n",
        "        self.tasks = tasks or {}\n",
        "        self.classifiers = nn.ModuleDict({\n",
        "            task: nn.Linear(self.model.config.hidden_size, num_labels) for task, num_labels in self.tasks.items()\n",
        "        })\n",
        "\n",
        "    def forward(self, input_ids, attention_mask, task, labels=None):\n",
        "        # Pass through the transformer with adapters\n",
        "        outputs = self.model(input_ids, attention_mask=attention_mask)\n",
        "        hidden_states = outputs.last_hidden_state\n",
        "        logits = self.classifiers[task](hidden_states[:, 0, :])  # CLS token\n",
        "        loss = None\n",
        "        if labels is not None:\n",
        "            loss_fct = nn.CrossEntropyLoss()\n",
        "            loss = loss_fct(logits.view(-1, self.classifiers[task].out_features), labels.view(-1))\n",
        "        return loss, logits\n",
        "\n",
        "    def add_task_tokens(self, texts, task):\n",
        "        # Add task-specific tokens to text\n",
        "        task_texts = [f\"[TASK-{task}] {text}\" for text in texts]\n",
        "        return self.encode_text(task_texts)\n",
        "\n",
        "# Train the multitask model with adapters, TensorBoard logging, and learning rate scheduler\n",
        "def train_with_scheduler(model, train_data, epochs=5, batch_size=32, learning_rate=5e-5, log_dir=\"./logs\", num_warmup_steps=500, num_training_steps=10000):\n",
        "    optimizer = AdamW(model.parameters(), lr=learning_rate, weight_decay=0.01)  # Add weight decay\n",
        "    scheduler = get_scheduler(optimizer, num_warmup_steps, num_training_steps)\n",
        "    writer = SummaryWriter(log_dir=log_dir)\n",
        "\n",
        "    # Create checkpoints directory if it doesn't exist\n",
        "    checkpoint_dir = \"./checkpoints\"\n",
        "    if not os.path.exists(checkpoint_dir):\n",
        "        os.makedirs(checkpoint_dir)\n",
        "\n",
        "    model.train()\n",
        "    for epoch in range(epochs):\n",
        "        total_loss = 0\n",
        "        for task, task_data in train_data.items():\n",
        "            train_dataset = TextDataset(task_data, model.tokenizer, for_classification=True)\n",
        "            train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
        "\n",
        "            for batch_idx, batch in enumerate(train_dataloader):\n",
        "                optimizer.zero_grad()\n",
        "                input_ids, attention_mask, labels = batch\n",
        "                input_ids, attention_mask, labels = input_ids.to(device), attention_mask.to(device), labels.to(device)\n",
        "\n",
        "                loss, logits = model(input_ids, attention_mask, task, labels=labels)\n",
        "                total_loss += loss.item()\n",
        "\n",
        "                loss.backward()\n",
        "                optimizer.step()\n",
        "                scheduler.step()\n",
        "\n",
        "                # Log loss for each batch\n",
        "                writer.add_scalar(f\"Loss/train_{task}\", loss.item(), epoch * len(train_dataloader) + batch_idx)\n",
        "\n",
        "        # Print loss and some sample predictions\n",
        "        print(f\"Epoch [{epoch + 1}/{epochs}], Task: {task}, Loss: {total_loss / len(train_dataloader)}\")\n",
        "        print(f\"Sample predictions: {logits[:5].cpu().detach().numpy()}\")\n",
        "        print(f\"Actual labels: {labels[:5].cpu().numpy()}\")\n",
        "\n",
        "        # Save model checkpoint at each epoch\n",
        "        torch.save(model.state_dict(), f\"./checkpoints/model_epoch_{epoch+1}.pt\")\n",
        "\n",
        "    writer.close()\n",
        "\n",
        "# Scheduler for learning rate\n",
        "def get_scheduler(optimizer, num_warmup_steps, num_training_steps):\n",
        "    def lr_lambda(current_step):\n",
        "        if current_step < num_warmup_steps:\n",
        "            return float(current_step) / float(max(1, num_warmup_steps))\n",
        "        return max(0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps)))\n",
        "\n",
        "    return LambdaLR(optimizer, lr_lambda)\n",
        "\n",
        "# Evaluation function with metrics\n",
        "def evaluate_with_metrics(model, test_data, task, batch_size=32):\n",
        "    test_dataloader = DataLoader(test_data, batch_size=batch_size)\n",
        "    model.eval()\n",
        "    all_labels, all_preds = [], []\n",
        "\n",
        "    with torch.no_grad():\n",
        "        for batch in test_dataloader:\n",
        "            input_ids, attention_mask, labels = batch\n",
        "            input_ids, attention_mask, labels = input_ids.to(device), attention_mask.to(device), labels.to(device)\n",
        "\n",
        "            _, logits = model(input_ids, attention_mask, task)\n",
        "            predictions = torch.argmax(logits, dim=-1)\n",
        "            all_labels.extend(labels.cpu().numpy())\n",
        "            all_preds.extend(predictions.cpu().numpy())\n",
        "\n",
        "    precision = precision_score(all_labels, all_preds, average='weighted')\n",
        "    recall = recall_score(all_labels, all_preds, average='weighted')\n",
        "    f1 = f1_score(all_labels, all_preds, average='weighted')\n",
        "\n",
        "    print(f\"Precision: {precision:.2f}, Recall: {recall:.2f}, F1 Score: {f1:.2f}\")\n",
        "    return precision, recall, f1\n",
        "\n",
        "# Synonym replacement for data augmentation\n",
        "def synonym_replacement(text, n=2):\n",
        "    words = text.split()\n",
        "    new_words = words.copy()\n",
        "    random.shuffle(words)\n",
        "\n",
        "    num_replaced = 0\n",
        "    for word in words:\n",
        "        synonyms = wordnet.synsets(word)\n",
        "        if synonyms:\n",
        "            synonym = synonyms[0].lemmas()[0].name()\n",
        "            new_words = [synonym if w == word and num_replaced < n else w for w in new_words]\n",
        "            num_replaced += 1\n",
        "        if num_replaced >= n:\n",
        "            break\n",
        "\n",
        "    return \" \".join(new_words)\n",
        "\n",
        "# Ensure wordnet is downloaded\n",
        "import nltk\n",
        "nltk.download('wordnet')\n",
        "\n",
        "# Example usage\n",
        "# Assuming train_data and test_data are available as lists of dictionaries with \"text\" and \"label\" fields\n",
        "\n",
        "train_data = {\n",
        "    \"task1\": [{\"text\": \"example sentence for task 1\", \"label\": 0}],  # Replace with actual data\n",
        "    \"task2\": [{\"text\": \"example sentence for task 2\", \"label\": 1}]   # Replace with actual data\n",
        "}\n",
        "\n",
        "tasks = {\"task1\": 2, \"task2\": 2}  # Define tasks with number of labels for each\n",
        "\n",
        "# Initialize the multitask model with adapters\n",
        "multitask_model = MultiTaskAdapterFoundationModel(model_name=\"bert-base-uncased\", tasks=tasks).to(device)\n",
        "\n",
        "# Train the multitask model with scheduler and logging\n",
        "train_with_scheduler(multitask_model, train_data, epochs=5, batch_size=32, learning_rate=5e-5, num_warmup_steps=500, num_training_steps=10000)\n",
        "\n",
        "# Test data\n",
        "test_data_task1 = [{\"text\": \"example test sentence for task 1\", \"label\": 0}]  # Replace with actual data\n",
        "test_dataset_task1 = TextDataset(test_data_task1, multitask_model.tokenizer, for_classification=True)\n",
        "\n",
        "# Evaluate the multitask model on a specific task\n",
        "evaluate_with_metrics(multitask_model, test_dataset_task1, task=\"task1\")\n",
        "\n",
        "# Example of synonym replacement\n",
        "text = \"The quick brown fox jumps over the lazy dog.\"\n",
        "augmented_text = synonym_replacement(text)\n",
        "print(f\"Original: {text}\")\n",
        "print(f\"Augmented: {augmented_text}\")"
      ],
      "metadata": {
        "id": "eQ_mx0WjIdHI"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}