myfirstfinetuning_flant5_peft.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94dfa40c-986e-49af-b742-b448f046bf45",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, TrainingArguments, Trainer\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c38fa3e0-cc61-4309-b3e9-68dce1fa10da",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define your dataset class\n",
    "class SimpleDataset(Dataset):\n",
    "    def __init__(self, input_texts, target_texts, tokenizer, max_length):\n",
    "        self.input_texts = input_texts\n",
    "        self.target_texts = target_texts\n",
    "        self.tokenizer = tokenizer\n",
    "        self.max_length = max_length\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.input_texts)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        input_text = self.input_texts[idx]\n",
    "        target_text = self.target_texts[idx]\n",
    "\n",
    "        input_encodings = self.tokenizer(input_text, return_tensors=\"pt\", padding=\"max_length\", truncation=True, max_length=self.max_length)\n",
    "        target_encodings = self.tokenizer(target_text, return_tensors=\"pt\", padding=\"max_length\", truncation=True, max_length=self.max_length)\n",
    "\n",
    "        # Remove unnecessary keys from target_encodings\n",
    "        target_encodings[\"input_ids\"] = target_encodings.pop(\"input_ids\")\n",
    "\n",
    "        return {\n",
    "            \"input_ids\": input_encodings[\"input_ids\"].squeeze(),\n",
    "            \"attention_mask\": input_encodings[\"attention_mask\"].squeeze(),\n",
    "            \"labels\": target_encodings[\"input_ids\"].squeeze(),\n",
    "        }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f954f9d-aa37-4493-827a-4865feae2d6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Your input and target texts\n",
    "input_texts = [\n",
    "    \"Who is your boss?\", \n",
    "    \"Hello\",\n",
    "] * 50  # Repeat for more instances  \n",
    "\n",
    "target_texts = [\n",
    "    \"You are my boss.\",\n",
    "    \"What is your problem, human?\",\n",
    "] * 50\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e610047-904d-4649-b193-b142224ef880",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the tokenizer and model\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"google/flan-t5-base\")\n",
    "model = AutoModelForSeq2SeqLM.from_pretrained(\"google/flan-t5-base\")\n",
    "\n",
    "tokenizer.pad_token = tokenizer.eos_token \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b87ac2ef-c1e6-4d06-8d25-973d3c9507bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set a maximum sequence length for both input and output\n",
    "max_length = 20\n",
    "\n",
    "# Create the dataset\n",
    "dataset = SimpleDataset(input_texts, target_texts, tokenizer, max_length)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7dbef0d2-9c54-43db-82e7-4beac037a0c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# output_dir=\"./flan_t5_fine_tuned_peft\"\n",
    "train_dataset = dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c28d8e02-48b7-4a74-abc2-0cf14b8b24c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import LoraConfig, get_peft_model, TaskType\n",
    "\n",
    "lora_config = LoraConfig(\n",
    "    r=32, # Rank\n",
    "    lora_alpha=32,\n",
    "    target_modules=[\"q\", \"v\"],\n",
    "    lora_dropout=0.05,\n",
    "    bias=\"none\",\n",
    "    task_type=TaskType.SEQ_2_SEQ_LM # FLAN-T5\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fedf4606-9666-4218-9f09-226b5c3e2dcc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_number_of_trainable_model_parameters(model):\n",
    "    trainable_model_params = 0\n",
    "    all_model_params = 0\n",
    "    for _, param in model.named_parameters():\n",
    "        all_model_params += param.numel()\n",
    "        if param.requires_grad:\n",
    "            trainable_model_params += param.numel()\n",
    "    return f\"trainable model parameters: {trainable_model_params}\\nall model parameters: {all_model_params}\\npercentage of trainable model parameters: {100 * trainable_model_params / all_model_params:.2f}%\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a5c93aa-b9e6-4b1c-8519-22353106166b",
   "metadata": {},
   "outputs": [],
   "source": [
    "peft_model = get_peft_model(model, \n",
    "                            lora_config)\n",
    "print(print_number_of_trainable_model_parameters(peft_model))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffdda183-8420-44d1-988e-9d8ec3d7ddfb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "output_dir = f'./flan_t5_fine_tuned_peft-{str(int(time.time()))}'\n",
    "\n",
    "peft_training_args = TrainingArguments(\n",
    "    output_dir=output_dir,\n",
    "    auto_find_batch_size=True,\n",
    "    learning_rate=1e-3, # Higher learning rate than full fine-tuning.\n",
    "    num_train_epochs=5,\n",
    "    logging_steps=100,\n",
    "    # max_steps=1    \n",
    ")\n",
    "    \n",
    "peft_trainer = Trainer(\n",
    "    model=peft_model,\n",
    "    args=peft_training_args,\n",
    "    train_dataset=train_dataset,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4109a97b-5d29-42c7-bb89-831ba6085442",
   "metadata": {},
   "outputs": [],
   "source": [
    "peft_trainer.train()\n",
    "\n",
    "peft_model_path=\"./flan_t5_fine_tuned_peft-local\"\n",
    "\n",
    "peft_trainer.model.save_pretrained(peft_model_path)\n",
    "tokenizer.save_pretrained(peft_model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7431f140-06e5-47cc-b70d-24ae28934cf6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForSeq2SeqLM, AutoTokenizer\n",
    "from peft import PeftModel, PeftConfig\n",
    "\n",
    "peft_model_base = AutoModelForSeq2SeqLM.from_pretrained(\"google/flan-t5-base\", torch_dtype=torch.bfloat16)\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"google/flan-t5-base\")\n",
    "\n",
    "peft_model = PeftModel.from_pretrained(model, \n",
    "                                       './flan_t5_fine_tuned_peft-local', \n",
    "                                       torch_dtype=torch.bfloat16,\n",
    "                                       is_trainable=False)\n",
    "print(print_number_of_trainable_model_parameters(peft_model))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85243e8e-ff57-4331-9f2f-265b12ae856f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}