lora.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyM9hQiYN/f+N9iqhs2OQyJJ",
      "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/jaebooker/low_rank_adapter_finetuning/blob/main/lora.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": "aEeGPudlyEVZ"
      },
      "outputs": [],
      "source": [
        "!pip3 install peft\n",
        "!pip3 install datasets\n",
        "from transformers import AutoModelForSeq2SeqLM\n",
        "from peft import PeftModel, PeftConfig\n",
        "import torch\n",
        "import os\n",
        "from transformers import AutoTokenizer\n",
        "from torch.utils.data import DataLoader\n",
        "from transformers import default_data_collator, get_linear_schedule_with_warmup\n",
        "from tqdm import tqdm\n",
        "from datasets import load_dataset\n",
        "\n",
        "dataset_name = \"twitter_complaints\"\n",
        "text_column = \"Tweet text\"\n",
        "label_column = \"text_label\"\n",
        "batch_size = 8\n",
        "\n",
        "peft_model_id = \"smangrul/twitter_complaints_bigscience_T0_3B_LORA_SEQ_2_SEQ_LM\"\n",
        "config = PeftConfig.from_pretrained(peft_model_id)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "peft_model_id = \"smangrul/twitter_complaints_bigscience_T0_3B_LORA_SEQ_2_SEQ_LM\"\n",
        "max_memory = {0: \"6GIB\", 1: \"0GIB\", 2: \"0GIB\", 3: \"0GIB\", 4: \"0GIB\", \"cpu\": \"30GB\"}\n",
        "config = PeftConfig.from_pretrained(peft_model_id)\n",
        "model = AutoModelForSeq2SeqLM.from_pretrained(config.base_model_name_or_path, device_map=\"auto\", max_memory=max_memory)\n",
        "model = PeftModel.from_pretrained(model, peft_model_id, device_map=\"auto\", max_memory=max_memory)"
      ],
      "metadata": {
        "id": "c3gW3w8BzOeY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dataset = load_dataset(\"ought/raft\", dataset_name)\n",
        "\n",
        "classes = [k.replace(\"_\", \" \") for k in dataset[\"train\"].features[\"Label\"].names]\n",
        "print(classes)\n",
        "dataset = dataset.map(\n",
        "    lambda x: {\"text_label\": [classes[label] for label in x[\"Label\"]]},\n",
        "    batched=True,\n",
        "    num_proc=1,\n",
        ")\n",
        "print(dataset)\n",
        "dataset[\"train\"][0]"
      ],
      "metadata": {
        "id": "e2TeoonnzxXP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)\n",
        "target_max_length = max([len(tokenizer(class_label)[\"input_ids\"]) for class_label in classes])\n",
        "\n",
        "\n",
        "def preprocess_function(examples):\n",
        "    inputs = examples[text_column]\n",
        "    targets = examples[label_column]\n",
        "    model_inputs = tokenizer(inputs, truncation=True)\n",
        "    labels = tokenizer(\n",
        "        targets, max_length=target_max_length, padding=\"max_length\", truncation=True, return_tensors=\"pt\"\n",
        "    )\n",
        "    labels = labels[\"input_ids\"]\n",
        "    labels[labels == tokenizer.pad_token_id] = -100\n",
        "    model_inputs[\"labels\"] = labels\n",
        "    return model_inputs\n",
        "\n",
        "\n",
        "processed_datasets = dataset.map(\n",
        "    preprocess_function,\n",
        "    batched=True,\n",
        "    num_proc=1,\n",
        "    remove_columns=dataset[\"train\"].column_names,\n",
        "    load_from_cache_file=True,\n",
        "    desc=\"Running tokenizer on dataset\",\n",
        ")\n",
        "\n",
        "train_dataset = processed_datasets[\"train\"]\n",
        "eval_dataset = processed_datasets[\"train\"]\n",
        "test_dataset = processed_datasets[\"test\"]\n",
        "\n",
        "\n",
        "def collate_fn(examples):\n",
        "    return tokenizer.pad(examples, padding=\"longest\", return_tensors=\"pt\")\n",
        "\n",
        "\n",
        "train_dataloader = DataLoader(\n",
        "    train_dataset, shuffle=True, collate_fn=collate_fn, batch_size=batch_size, pin_memory=True\n",
        ")\n",
        "eval_dataloader = DataLoader(eval_dataset, collate_fn=collate_fn, batch_size=batch_size, pin_memory=True)\n",
        "test_dataloader = DataLoader(test_dataset, collate_fn=collate_fn, batch_size=batch_size, pin_memory=True)"
      ],
      "metadata": {
        "id": "NPX-VaJ_0J3x"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model.eval()\n",
        "i = 15\n",
        "inputs = tokenizer(f'{text_column} : {dataset[\"test\"][i][\"Tweet text\"]} Label : ', return_tensors=\"pt\")\n",
        "print(dataset[\"test\"][i][\"Tweet text\"])\n",
        "print(inputs)\n",
        "\n",
        "with torch.no_grad():\n",
        "    outputs = model.generate(input_ids=inputs[\"input_ids\"].to(\"cuda\"), max_new_tokens=10)\n",
        "    print(outputs)\n",
        "    print(tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True))"
      ],
      "metadata": {
        "id": "QQomFjZ-0PI2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model.eval()\n",
        "eval_preds = []\n",
        "for _, batch in enumerate(tqdm(eval_dataloader)):\n",
        "    batch = {k: v.to(\"cuda\") for k, v in batch.items() if k != \"labels\"}\n",
        "    with torch.no_grad():\n",
        "        outputs = model.generate(**batch, max_new_tokens=10)\n",
        "    preds = outputs.detach().cpu().numpy()\n",
        "    eval_preds.extend(tokenizer.batch_decode(preds, skip_special_tokens=True))"
      ],
      "metadata": {
        "id": "Crhoc1yf0S08"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "correct = 0\n",
        "total = 0\n",
        "for pred, true in zip(eval_preds, dataset[\"train\"][label_column]):\n",
        "    if pred.strip() == true.strip():\n",
        "        correct += 1\n",
        "    total += 1\n",
        "accuracy = correct / total * 100\n",
        "print(f\"{accuracy=}\")\n",
        "print(f\"{eval_preds[:10]=}\")\n",
        "print(f\"{dataset['train'][label_column][:10]=}\")"
      ],
      "metadata": {
        "id": "Qw5vq6T50V4W"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model.eval()\n",
        "test_preds = []\n",
        "\n",
        "for _, batch in enumerate(tqdm(test_dataloader)):\n",
        "    batch = {k: v for k, v in batch.items() if k != \"labels\"}\n",
        "    with torch.no_grad():\n",
        "        outputs = model.generate(**batch, max_new_tokens=10)\n",
        "    preds = outputs.detach().cpu().numpy()\n",
        "    test_preds.extend(tokenizer.batch_decode(preds, skip_special_tokens=True))\n",
        "    if len(test_preds) > 100:\n",
        "        break\n",
        "test_preds"
      ],
      "metadata": {
        "id": "guo_yaZU0ZFE"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}