paligemma_finetuning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export HF_TOKEN = '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'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from datasets import load_dataset, load_from_disk\n",
    "from transformers import PaliGemmaProcessor, PaliGemmaForConditionalGeneration, BitsAndBytesConfig, TrainingArguments, Trainer\n",
    "import torch\n",
    "from peft import get_peft_model, LoraConfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"HF_TOKEN\"] =  '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'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Loading the DataSet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\dhanu\\.conda\\envs\\llm\\lib\\site-packages\\datasets\\load.py:1461: FutureWarning: The repository for HuggingFaceM4/VQAv2 contains custom code which must be executed to correctly load the dataset. You can inspect the repository content at https://hf.co/datasets/HuggingFaceM4/VQAv2\n",
      "You can avoid this message in future by passing the argument `trust_remote_code=True`.\n",
      "Passing `trust_remote_code=True` will be mandatory to load this dataset from the next major release of `datasets`.\n",
      "  warnings.warn(\n",
      "Repo card metadata block was not found. Setting CardData to empty.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'multiple_choice_answer': 'yes', 'question': 'Is it probably cold here?', 'image': <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=640x501 at 0x28168E65930>}\n"
     ]
    }
   ],
   "source": [
    "data = load_dataset('HuggingFaceM4/VQAv2', split=\"train[:10%]\")\n",
    "cols_remove = [\"question_type\", \"answers\", \"answer_type\", \"image_id\", \"question_id\"]\n",
    "data = data.remove_columns(cols_remove)\n",
    "split_data = data.train_test_split(test_size=0.05)\n",
    "train_data = split_data[\"test\"]\n",
    "print(train_data[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Model "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`config.hidden_act` is ignored, you should use `config.hidden_activation` instead.\n",
      "Gemma's activation function will be set to `gelu_pytorch_tanh`. Please, use\n",
      "`config.hidden_activation` if you want to override this behaviour.\n",
      "See https://github.com/huggingface/transformers/pull/29402 for more details.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "43eea246f591454eb89cbb855d147e71",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "model_id = \"leo009/paligemma-3b-pt-224\"\n",
    "#model_id = \"google/paligemma-3b-pt-224\"\n",
    "processor = PaliGemmaProcessor.from_pretrained(model_id)\n",
    "device = \"cuda\"\n",
    "image_token = processor.tokenizer.convert_tokens_to_ids(\"<image>\")\n",
    "model = PaliGemmaForConditionalGeneration.from_pretrained(model_id, torch_dtype=torch.bfloat16).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "for param in model.vision_tower.parameters():\n",
    "    param.required_grad = True \n",
    "\n",
    "for param in model.multi_modal_projector.parameters():\n",
    "    param.required_grad = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "True\n",
      "True\n",
      "False\n",
      "False\n",
      "False\n"
     ]
    }
   ],
   "source": [
    "for param in model.parameters():\n",
    "    print(param.requires_grad)  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Loading Quantised Model "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit = True,\n",
    "    bnb_4bit_quant_type = \"nf4\",\n",
    "    bnb_4bit_compute_dtype= torch.bfloat16\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "lora_config = LoraConfig(\n",
    "    r = 8,\n",
    "    lora_alpha = 32,\n",
    "    lora_dropout = 0.05,\n",
    "    task_type = \"CAUSAL_LM\",\n",
    "    target_modules=[\"q_proj\", \"o_proj\", \"k_proj\", \"v_proj\",\n",
    "                    \"gate_proj\", \"up_proj\", \"down_proj\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "570fbfb42fce4b0c97c3b2f820f8e57e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 11,298,816 || all params: 2,934,765,296 || trainable%: 0.3849989644964099\n"
     ]
    }
   ],
   "source": [
    "model = PaliGemmaForConditionalGeneration.from_pretrained(model_id,quantization_config=bnb_config, device_map = {\"\":0})\n",
    "\n",
    "model = get_peft_model(model,lora_config)\n",
    "\n",
    "model.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Finetuning the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def collate_fn(examples):\n",
    "    texts = [\"answer \" + example[\"question\"] for example in examples]\n",
    "    labels = [example['multiple_choice_answer'] for example in examples]\n",
    "    images = [example[\"image\"].convert(\"RGB\") for example in examples]\n",
    "    tokens = processor(text=texts, images=images, suffix=labels,\n",
    "                       return_tensors=\"pt\", padding=\"longest\",\n",
    "                       tokenize_newline_separately=False)\n",
    "    tokens = tokens.to(torch.bfloat16).to(device)\n",
    "    return tokens "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import TrainingArguments\n",
    "args=TrainingArguments(\n",
    "            num_train_epochs=2,\n",
    "            remove_unused_columns=False,\n",
    "            per_device_train_batch_size=4,\n",
    "            gradient_accumulation_steps=4,\n",
    "            warmup_steps=2,\n",
    "            learning_rate=2e-5,\n",
    "            weight_decay=1e-6,\n",
    "            adam_beta2=0.999,\n",
    "            logging_steps=100,\n",
    "            optim=\"adamw_hf\",\n",
    "            save_strategy=\"steps\",\n",
    "            save_steps=1000,\n",
    "            push_to_hub=True,\n",
    "            save_total_limit=1,\n",
    "            output_dir=\"paligemma_vqav2\",\n",
    "            bf16=True,\n",
    "            report_to=[\"tensorboard\"],\n",
    "            dataloader_pin_memory=False\n",
    "        )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model = model,\n",
    "    train_dataset = train_data,\n",
    "    data_collator=collate_fn,\n",
    "    args = args\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "torch.cuda.empty_cache()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6ad743a96e8d4587a88a596acd2399be",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/276 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 49391.23, 'grad_norm': nan, 'learning_rate': 2e-05, 'epoch': 0.72}\n",
      "{'loss': 0.0, 'grad_norm': nan, 'learning_rate': 2e-05, 'epoch': 1.44}\n",
      "{'train_runtime': 27857.5245, 'train_samples_per_second': 0.159, 'train_steps_per_second': 0.01, 'train_loss': 17895.373188405796, 'epoch': 1.99}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=276, training_loss=17895.373188405796, metrics={'train_runtime': 27857.5245, 'train_samples_per_second': 0.159, 'train_steps_per_second': 0.01, 'total_flos': 1.7286077798096352e+16, 'train_loss': 17895.373188405796, 'epoch': 1.9891891891891893})"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/Dhanushkumar/paligemma_vqav2/commit/a1c0517d4345e085ae2793bea386f0606ccceadf', commit_message='Dhanushkumar/paligemma_VQAv2', commit_description='', oid='a1c0517d4345e085ae2793bea386f0606ccceadf', pr_url=None, pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.push_to_hub(\"Dhanushkumar/paligemma_VQAv2\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test the trained Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import site\n",
    "from peft import PeftModel, PeftConfig\n",
    "from transformers import AutoModelForCausalLM, PaliGemmaProcessor, AutoModelForPreTraining\n",
    "from huggingface_hub import login\n",
    "import torch\n",
    "from PIL import Image\n",
    "import requests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup site packages path\n",
    "python_version = f\"{sys.version_info.major}.{sys.version_info.minor}\"\n",
    "site_packages_path = os.path.expanduser(f'~/.local/lib/python{python_version}/site-packages')\n",
    "site.addsitedir(site_packages_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The token has not been saved to the git credentials helper. Pass `add_to_git_credential=True` in this function directly or `--add-to-git-credential` if using via `huggingface-cli` if you want to set the git credential as well.\n",
      "Token is valid (permission: write).\n",
      "Your token has been saved to C:\\Users\\dhanu\\.cache\\huggingface\\token\n",
      "Login successful\n"
     ]
    }
   ],
   "source": [
    "# Login to Hugging Face Hub\n",
    "token = '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'\n",
    "login(token=token)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`Load PeftConfig and Base Model`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = PeftConfig.from_pretrained(\"Dhanushkumar/paligemma_VQAv2\")\n",
    "\n",
    "base_model = AutoModelForPreTraining.from_pretrained(\"leo009/paligemma-3b-pt-224\")\n",
    "model = PeftModel.from_pretrained(base_model,\"Dhanushkumar/paligemma_VQAv2\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Loading and Processing the Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_text = \"What is in this image?\"\n",
    "img_url = \"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/bee.JPG?download=true\"\n",
    "input_image = Image.open(requests.get(img_url, stream=True).raw)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loading PaliGemma Processor\n",
    "processor = PaliGemmaProcessor.from_pretrained(\"leo009/paligemma-3b-pt-224\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Preprocessing Inputs\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "inputs = processor(text=input_text, images=input_image, padding=\"longest\", do_convert_rgb=True, return_tensors=\"pt\").to(device)\n",
    "model.to(device)\n",
    "inputs = inputs.to(dtype=model.dtype)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generating and Decoding Output\n",
    "with torch.no_grad():\n",
    "    output = model.generate(**inputs, max_length=496)\n",
    "\n",
    "print(processor.decode(output[0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Analyse VQAv2 Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import re\n",
    "from datasets import load_dataset\n",
    "from PIL import Image\n",
    "\n",
    "# Create a directory to save images\n",
    "os.makedirs('saved_images', exist_ok=True)\n",
    "\n",
    "# Load and preprocess the dataset\n",
    "ds = load_dataset('HuggingFaceM4/VQAv2', split=\"train[:10%]\")\n",
    "cols_keep = [\"question_type\", \"answers\", \"answer_type\", \"image\", \"image_id\", \"question_id\", \"question\"]\n",
    "ds = ds.remove_columns([col for col in ds.column_names if col not in cols_keep])\n",
    "split_ds = ds.train_test_split(test_size=0.05)\n",
    "train_ds = split_ds[\"test\"]\n",
    "\n",
    "# Helper function to sanitize filenames\n",
    "def sanitize_filename(text):\n",
    "    text = re.sub(r'\\W+', '_', text)\n",
    "    return text[:50]  # Truncate to 50 characters for safety\n",
    "\n",
    "# Save 10 images from the dataset with filenames based on metadata\n",
    "for i in range(10):\n",
    "    example = train_ds[i]\n",
    "    image = example[\"image\"]\n",
    "    question_type = sanitize_filename(example[\"question_type\"])\n",
    "    question = sanitize_filename(example[\"question\"])\n",
    "    # Combine unique answers, truncating and sanitizing the string\n",
    "    answers = \"_\".join(sorted(set([sanitize_filename(ans['answer']) for ans in example[\"answers\"]])))\n",
    "    answers = answers[:50]  # Truncate to 50 characters for safety\n",
    "    answer_type = sanitize_filename(example[\"answer_type\"])\n",
    "    \n",
    "    filename = f\"{question_type}_{question}_{answers}_{answer_type}.jpg\"\n",
    "    filepath = os.path.join('saved_images', filename)\n",
    "    \n",
    "    image.save(filepath)\n",
    "\n",
    "print(\"10 images saved successfully.\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llm",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}