Custom_Learning_Rate_Schedulers.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "authorship_tag": "ABX9TyNl/s5Ov8MBZyF2z6HQL9Qw",
      "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/Custom_Learning_Rate_Schedulers.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HF6Np15y20NH"
      },
      "outputs": [],
      "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 the FoundationModel class\n",
        "class FoundationModel(nn.Module):\n",
        "    def __init__(self, model_name=\"bert-base-uncased\"):\n",
        "        super(FoundationModel, self).__init__()\n",
        "        self.model = BertModel.from_pretrained(model_name)\n",
        "        self.tokenizer = BertTokenizer.from_pretrained(model_name)\n",
        "\n",
        "    def forward(self, input_ids, attention_mask):\n",
        "        outputs = self.model(input_ids, attention_mask=attention_mask)\n",
        "        return 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 the MultiTaskFoundationModel class for multitask learning\n",
        "class MultiTaskFoundationModel(FoundationModel):\n",
        "    def __init__(self, model_name=\"bert-base-uncased\", tasks=None):\n",
        "        super().__init__(model_name)\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\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 TensorBoard logging and learning rate scheduler\n",
        "def train_with_scheduler(model, train_data, epochs=3, 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)\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, _ = 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",
        "        # Save model checkpoint at each epoch\n",
        "        torch.save(model.state_dict(), f\"./checkpoints/model_epoch_{epoch+1}.pt\")\n",
        "        print(f\"Epoch [{epoch + 1}/{epochs}], Task: {task}, Loss: {total_loss / len(train_dataloader)}\")\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\n",
        "multitask_model = MultiTaskFoundationModel(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)\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}\")"
      ]
    }
  ]
}