{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "63b957e0-d83b-48a6-8ae0-a276b983e181",
   "metadata": {},
   "source": [
    "### Optional: install the necessary packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5569a330-ea6b-4402-9fd5-e7a0ce981bc9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: huggingface_hub in /workspace/.miniconda3/lib/python3.12/site-packages (0.26.2)\n",
      "Requirement already satisfied: filelock in /workspace/.miniconda3/lib/python3.12/site-packages (from huggingface_hub) (3.16.1)\n",
      "Requirement already satisfied: fsspec>=2023.5.0 in /workspace/.miniconda3/lib/python3.12/site-packages (from huggingface_hub) (2024.9.0)\n",
      "Requirement already satisfied: packaging>=20.9 in /workspace/.miniconda3/lib/python3.12/site-packages (from huggingface_hub) (24.1)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /workspace/.miniconda3/lib/python3.12/site-packages (from huggingface_hub) (6.0.2)\n",
      "Requirement already satisfied: requests in /workspace/.miniconda3/lib/python3.12/site-packages (from huggingface_hub) (2.32.3)\n",
      "Requirement already satisfied: tqdm>=4.42.1 in /workspace/.miniconda3/lib/python3.12/site-packages (from huggingface_hub) (4.66.5)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /workspace/.miniconda3/lib/python3.12/site-packages (from huggingface_hub) (4.11.0)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /workspace/.miniconda3/lib/python3.12/site-packages (from requests->huggingface_hub) (3.3.2)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /workspace/.miniconda3/lib/python3.12/site-packages (from requests->huggingface_hub) (3.7)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /workspace/.miniconda3/lib/python3.12/site-packages (from requests->huggingface_hub) (2.2.3)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /workspace/.miniconda3/lib/python3.12/site-packages (from requests->huggingface_hub) (2024.8.30)\n",
      "Note: you may need to restart the kernel to use updated packages.\n",
      "Requirement already satisfied: datasets in /workspace/.miniconda3/lib/python3.12/site-packages (3.1.0)\n",
      "Requirement already satisfied: filelock in /workspace/.miniconda3/lib/python3.12/site-packages (from datasets) (3.16.1)\n",
      "Requirement already satisfied: numpy>=1.17 in /workspace/.miniconda3/lib/python3.12/site-packages (from datasets) (2.1.3)\n",
      "Requirement already satisfied: pyarrow>=15.0.0 in /workspace/.miniconda3/lib/python3.12/site-packages (from datasets) (18.0.0)\n",
      "Requirement already satisfied: dill<0.3.9,>=0.3.0 in /workspace/.miniconda3/lib/python3.12/site-packages (from datasets) (0.3.8)\n",
      "Requirement already satisfied: pandas in /workspace/.miniconda3/lib/python3.12/site-packages (from datasets) (2.2.3)\n",
      "Requirement already satisfied: requests>=2.32.2 in /workspace/.miniconda3/lib/python3.12/site-packages (from datasets) (2.32.3)\n",
      "Requirement already satisfied: tqdm>=4.66.3 in /workspace/.miniconda3/lib/python3.12/site-packages (from datasets) (4.66.5)\n",
      "Requirement already satisfied: xxhash in /workspace/.miniconda3/lib/python3.12/site-packages (from datasets) (3.5.0)\n",
      "Requirement already satisfied: multiprocess<0.70.17 in /workspace/.miniconda3/lib/python3.12/site-packages (from datasets) (0.70.16)\n",
      "Requirement already satisfied: fsspec<=2024.9.0,>=2023.1.0 in /workspace/.miniconda3/lib/python3.12/site-packages (from fsspec[http]<=2024.9.0,>=2023.1.0->datasets) (2024.9.0)\n",
      "Requirement already satisfied: aiohttp in /workspace/.miniconda3/lib/python3.12/site-packages (from datasets) (3.11.2)\n",
      "Requirement already satisfied: huggingface-hub>=0.23.0 in /workspace/.miniconda3/lib/python3.12/site-packages (from datasets) (0.26.2)\n",
      "Requirement already satisfied: packaging in /workspace/.miniconda3/lib/python3.12/site-packages (from datasets) (24.1)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /workspace/.miniconda3/lib/python3.12/site-packages (from datasets) (6.0.2)\n",
      "Requirement already satisfied: aiohappyeyeballs>=2.3.0 in /workspace/.miniconda3/lib/python3.12/site-packages (from aiohttp->datasets) (2.4.3)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /workspace/.miniconda3/lib/python3.12/site-packages (from aiohttp->datasets) (1.3.1)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /workspace/.miniconda3/lib/python3.12/site-packages (from aiohttp->datasets) (24.2.0)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /workspace/.miniconda3/lib/python3.12/site-packages (from aiohttp->datasets) (1.5.0)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /workspace/.miniconda3/lib/python3.12/site-packages (from aiohttp->datasets) (6.1.0)\n",
      "Requirement already satisfied: propcache>=0.2.0 in /workspace/.miniconda3/lib/python3.12/site-packages (from aiohttp->datasets) (0.2.0)\n",
      "Requirement already satisfied: yarl<2.0,>=1.17.0 in /workspace/.miniconda3/lib/python3.12/site-packages (from aiohttp->datasets) (1.17.1)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /workspace/.miniconda3/lib/python3.12/site-packages (from huggingface-hub>=0.23.0->datasets) (4.11.0)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /workspace/.miniconda3/lib/python3.12/site-packages (from requests>=2.32.2->datasets) (3.3.2)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /workspace/.miniconda3/lib/python3.12/site-packages (from requests>=2.32.2->datasets) (3.7)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /workspace/.miniconda3/lib/python3.12/site-packages (from requests>=2.32.2->datasets) (2.2.3)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /workspace/.miniconda3/lib/python3.12/site-packages (from requests>=2.32.2->datasets) (2024.8.30)\n",
      "Requirement already satisfied: python-dateutil>=2.8.2 in /workspace/.miniconda3/lib/python3.12/site-packages (from pandas->datasets) (2.9.0.post0)\n",
      "Requirement already satisfied: pytz>=2020.1 in /workspace/.miniconda3/lib/python3.12/site-packages (from pandas->datasets) (2024.1)\n",
      "Requirement already satisfied: tzdata>=2022.7 in /workspace/.miniconda3/lib/python3.12/site-packages (from pandas->datasets) (2024.2)\n",
      "Requirement already satisfied: six>=1.5 in /workspace/.miniconda3/lib/python3.12/site-packages (from python-dateutil>=2.8.2->pandas->datasets) (1.16.0)\n",
      "Note: you may need to restart the kernel to use updated packages.\n",
      "Requirement already satisfied: bitsandbytes in /workspace/.miniconda3/lib/python3.12/site-packages (0.44.2.dev0)\n",
      "Requirement already satisfied: torch in /workspace/.miniconda3/lib/python3.12/site-packages (from bitsandbytes) (2.5.1)\n",
      "Requirement already satisfied: numpy in /workspace/.miniconda3/lib/python3.12/site-packages (from bitsandbytes) (2.1.3)\n",
      "Requirement already satisfied: filelock in /workspace/.miniconda3/lib/python3.12/site-packages (from torch->bitsandbytes) (3.16.1)\n",
      "Requirement already satisfied: typing-extensions>=4.8.0 in /workspace/.miniconda3/lib/python3.12/site-packages (from torch->bitsandbytes) (4.11.0)\n",
      "Requirement already satisfied: networkx in /workspace/.miniconda3/lib/python3.12/site-packages (from torch->bitsandbytes) (3.4.2)\n",
      "Requirement already satisfied: jinja2 in /workspace/.miniconda3/lib/python3.12/site-packages (from torch->bitsandbytes) (3.1.4)\n",
      "Requirement already satisfied: fsspec in /workspace/.miniconda3/lib/python3.12/site-packages (from torch->bitsandbytes) (2024.9.0)\n",
      "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.4.127 in /workspace/.miniconda3/lib/python3.12/site-packages (from torch->bitsandbytes) (12.4.127)\n",
      "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.4.127 in /workspace/.miniconda3/lib/python3.12/site-packages (from torch->bitsandbytes) (12.4.127)\n",
      "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.4.127 in /workspace/.miniconda3/lib/python3.12/site-packages (from torch->bitsandbytes) (12.4.127)\n",
      "Requirement already satisfied: nvidia-cudnn-cu12==9.1.0.70 in /workspace/.miniconda3/lib/python3.12/site-packages (from torch->bitsandbytes) (9.1.0.70)\n",
      "Requirement already satisfied: nvidia-cublas-cu12==12.4.5.8 in /workspace/.miniconda3/lib/python3.12/site-packages (from torch->bitsandbytes) (12.4.5.8)\n",
      "Requirement already satisfied: nvidia-cufft-cu12==11.2.1.3 in /workspace/.miniconda3/lib/python3.12/site-packages (from torch->bitsandbytes) (11.2.1.3)\n",
      "Requirement already satisfied: nvidia-curand-cu12==10.3.5.147 in /workspace/.miniconda3/lib/python3.12/site-packages (from torch->bitsandbytes) (10.3.5.147)\n",
      "Requirement already satisfied: nvidia-cusolver-cu12==11.6.1.9 in /workspace/.miniconda3/lib/python3.12/site-packages (from torch->bitsandbytes) (11.6.1.9)\n",
      "Requirement already satisfied: nvidia-cusparse-cu12==12.3.1.170 in /workspace/.miniconda3/lib/python3.12/site-packages (from torch->bitsandbytes) (12.3.1.170)\n",
      "Requirement already satisfied: nvidia-nccl-cu12==2.21.5 in /workspace/.miniconda3/lib/python3.12/site-packages (from torch->bitsandbytes) (2.21.5)\n",
      "Requirement already satisfied: nvidia-nvtx-cu12==12.4.127 in /workspace/.miniconda3/lib/python3.12/site-packages (from torch->bitsandbytes) (12.4.127)\n",
      "Requirement already satisfied: nvidia-nvjitlink-cu12==12.4.127 in /workspace/.miniconda3/lib/python3.12/site-packages (from torch->bitsandbytes) (12.4.127)\n",
      "Requirement already satisfied: triton==3.1.0 in /workspace/.miniconda3/lib/python3.12/site-packages (from torch->bitsandbytes) (3.1.0)\n",
      "Requirement already satisfied: setuptools in /workspace/.miniconda3/lib/python3.12/site-packages (from torch->bitsandbytes) (75.1.0)\n",
      "Requirement already satisfied: sympy==1.13.1 in /workspace/.miniconda3/lib/python3.12/site-packages (from torch->bitsandbytes) (1.13.1)\n",
      "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /workspace/.miniconda3/lib/python3.12/site-packages (from sympy==1.13.1->torch->bitsandbytes) (1.3.0)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /workspace/.miniconda3/lib/python3.12/site-packages (from jinja2->torch->bitsandbytes) (2.1.3)\n",
      "Note: you may need to restart the kernel to use updated packages.\n",
      "Note: you may need to restart the kernel to use updated packages.\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "!git config --global credential.helper store\n",
    "%pip install huggingface_hub\n",
    "%pip install -U datasets\n",
    "%pip install -U bitsandbytes\n",
    "%pip install -q git+https://github.com/huggingface/transformers.git\n",
    "%pip install -q accelerate datasets peft torchvision torchaudio"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca7f5161-0104-45c4-83c9-1dd0dad15e29",
   "metadata": {},
   "source": [
    "## Login on Hugging Face"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "14c445aa-c8bc-43b9-8d18-46d79055e1f0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hugging Face token found in environment variable\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Note: Environment variable`HF_TOKEN` is set and is the current active token independently from the token you've just configured.\n"
     ]
    }
   ],
   "source": [
    "from huggingface_hub import login\n",
    "import os\n",
    "\n",
    "HF_TOKEN = \"hf_C…………\"\n",
    "\n",
    "if os.environ.get('HF_TOKEN') is not None:\n",
    "  HF_TOKEN = os.environ.get('HF_TOKEN')\n",
    "  print(f\"Hugging Face token found in environment variable\")\n",
    "try:\n",
    "  import google.colab\n",
    "  from google.colab import userdata\n",
    "  if (userdata.get('HF_TOKEN') is not None) and (HF_TOKEN == \"\"):\n",
    "    HF_TOKEN = userdata.get('HF_TOKEN')\n",
    "  else:\n",
    "    raise ValueError(\"Please set your Hugging Face token in the user data panel, or pass it as an environment variable\")\n",
    "except ModuleNotFoundError:\n",
    "  if HF_TOKEN is None:\n",
    "    raise ValueError(\"Please set your Hugging Face token in the user data panel, or pass it as an environment variable\")\n",
    "\n",
    "login(\n",
    "  token=HF_TOKEN,\n",
    "  add_to_git_credential=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c33a0325-f056-43c4-a500-6ad1b5632ee5",
   "metadata": {},
   "source": [
    "### Set the environment variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bfd9dbb4-59c4-4d97-9727-5c2d9e60724e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#source_model_id = \"HuggingFaceM4/Idefics3-8B-Llama3\"\n",
    "source_model_id = \"meta-llama/Llama-3.2-11B-Vision-Instruct\"\n",
    "detination_model_id = \"eltorio/IDEFICS3_medical_instruct\"\n",
    "dataset_id = \"ruslanmv/ai-medical-dataset\"\n",
    "prompt= \"You are a medical doctor with 15 year of experience verifying the knowledge of a new diploma medical doctor\"\n",
    "output_dir = \"IDEFICS3_medical_instruct\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da89fd1d-44dd-4d4e-803f-ebf58b23165f",
   "metadata": {},
   "source": [
    "### Optionally clone the model repository"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f6510468-c903-4bd2-80fd-2511b7fb2f72",
   "metadata": {},
   "outputs": [],
   "source": [
    "# clone Hugging Face model repository\n",
    "# !git clone https://huggingface.co/$destination_model_id $output_dir"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bfc04cd-3867-445f-9764-21bc72a07f60",
   "metadata": {},
   "source": [
    "### Load the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1efffd3f-aece-4858-b615-8fb1f2997068",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f67282ac10514645b9ffcdb9f797cf22",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Resolving data files:   0%|          | 0/18 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0ff3d2de672641c8b59a0518d2987e4f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Resolving data files:   0%|          | 0/18 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "48453247d4564242b8a5c216fc3d7fe7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading dataset shards:   0%|          | 0/18 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "base_dataset = load_dataset(\"ruslanmv/ai-medical-dataset\")\n",
    "# define the train dataset as a random 80% of the data\n",
    "train_dataset = base_dataset[\"train\"].train_test_split(test_size=0.2)[\"train\"]\n",
    "# define the eval dataset as the remaining 20%\n",
    "eval_dataset = base_dataset[\"train\"].train_test_split(test_size=0.2)[\"test\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6be720c0-90a3-4ec0-bf1c-4d18e6b2bb31",
   "metadata": {
    "id": "Tu199rI3AMtG"
   },
   "source": [
    "### Configure LoRA adapters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86163110-f084-429e-bc70-5b281a679d1c",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 67,
     "referenced_widgets": [
      "9f270a1faba44c1194d683a431d43f6f",
      "7403ea0ef7f04e2389c351b0e10a68ce",
      "fbc17ebe22854738b98bf2488c6105b1",
      "22e04cb812df4f5ca621b041c213d17e",
      "96d4f59ed68749eaaf9ad0c6a4e94f13",
      "7a8bc62ce92a4571ad42ac2beef851ef",
      "ae0117f952d94e469793bd8db3ad9a19",
      "0b8f8a93515d49488642b008e47b4bf3",
      "c9691e94467b4badbb71116136350c93",
      "fcd773b4e6154af78e7d476c59294e25",
      "7bd9e3fc12e64bc7bb95891dcd4153c7"
     ]
    },
    "executionInfo": {
     "elapsed": 106706,
     "status": "ok",
     "timestamp": 1730998316179,
     "user": {
      "displayName": "Ronan Le Meillat",
      "userId": "09161391957806824350"
     },
     "user_tz": -60
    },
    "id": "GTf6VxwwAMtG",
    "outputId": "86a93163-6b30-4e8b-ca1c-7392922a4aa6"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`low_cpu_mem_usage` was None, now default to True since model is quantized.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "932c25f63b9242f0b2bfb42004a8362b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "from peft import LoraConfig, get_peft_model\n",
    "from transformers import AutoProcessor, MllamaForConditionalGeneration, BitsAndBytesConfig, Idefics3ForConditionalGeneration\n",
    "\n",
    "DEVICE = \"cuda:0\"\n",
    "USE_LORA = False\n",
    "USE_QLORA = True\n",
    "\n",
    "processor = AutoProcessor.from_pretrained(\n",
    "    source_model_id,\n",
    "    do_image_splitting=False\n",
    ")\n",
    "\n",
    "if USE_QLORA or USE_LORA:\n",
    "    lora_config = LoraConfig(\n",
    "        r=8,\n",
    "        lora_alpha=8,\n",
    "        lora_dropout=0.1,\n",
    "        target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n",
    "                      \"gate_proj\", \"up_proj\", \"down_proj\",],\n",
    "        #target_modules='.*(text_model|modality_projection|perceiver_resampler).*(down_proj|gate_proj|up_proj|k_proj|q_proj|v_proj|o_proj).*$',\n",
    "        task_type=\"CAUSAL_LM\",\n",
    "        use_dora=False if USE_QLORA else True,\n",
    "        init_lora_weights=\"gaussian\"\n",
    "    )\n",
    "    if USE_QLORA:\n",
    "        bnb_config = BitsAndBytesConfig(\n",
    "            load_in_4bit=True,\n",
    "            bnb_4bit_quant_type=\"nf4\",\n",
    "            bnb_4bit_compute_dtype=torch.float16\n",
    "        )\n",
    "    model = MllamaForConditionalGeneration.from_pretrained(\n",
    "        source_model_id,\n",
    "        torch_dtype=torch.float16,\n",
    "        quantization_config=bnb_config if USE_QLORA else None,\n",
    "    )\n",
    "    model = get_peft_model(model, lora_config)\n",
    "    #model.add_adapter(lora_config)\n",
    "    #model.enable_adapters()\n",
    "else:\n",
    "    model = MllamaForConditionalGeneration.from_pretrained(\n",
    "        source_model_id,\n",
    "        torch_dtype=torch.float16,\n",
    "        _attn_implementation=\"flash_attention_2\", # This works for A100 or H100\n",
    "    ).to(DEVICE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffbbd54b-e579-44ef-9652-cd8496b2fd4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2c2fd34-e1e4-427d-86d0-73bf74ff0005",
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_dataset\n",
    "eval_dataset[2310]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01b0fe64-1f62-42a8-847d-9162f5015c4e",
   "metadata": {
    "id": "0JeaGZxHAMtG"
   },
   "source": [
    "### Create Data Collator for IDEFICS3 format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29d96aea-445d-482d-b7dc-861635a5389c",
   "metadata": {
    "executionInfo": {
     "elapsed": 426,
     "status": "ok",
     "timestamp": 1730998596513,
     "user": {
      "displayName": "Ronan Le Meillat",
      "userId": "09161391957806824350"
     },
     "user_tz": -60
    },
    "id": "X6TWyPHaAMtH"
   },
   "outputs": [],
   "source": [
    "class MyDataCollator:\n",
    "    def __init__(self, processor):\n",
    "        self.processor = processor\n",
    "        self.image_token_id = processor.tokenizer.additional_special_tokens_ids[\n",
    "            processor.tokenizer.additional_special_tokens.index(\"<image>\")\n",
    "        ]\n",
    "\n",
    "    def __call__(self, samples):\n",
    "        texts = []\n",
    "        images = []\n",
    "        for sample in samples:\n",
    "            question = sample[\"question\"]\n",
    "            answer = sample[\"context\"]\n",
    "            messages = [\n",
    "                {\n",
    "                    \"role\": \"system\",\n",
    "                    \"content\": [\n",
    "                        {\"type\": \"text\", \"text\": prompt}\n",
    "                    ]\n",
    "\n",
    "                },\n",
    "                {\n",
    "                    \"role\": \"user\",\n",
    "                    \"content\": [\n",
    "                        {\"type\": \"text\", \"text\": question },\n",
    "                    ]\n",
    "                },\n",
    "                {\n",
    "                    \"role\": \"assistant\",\n",
    "                    \"content\": [\n",
    "                        {\"type\": \"text\", \"text\": answer}\n",
    "                    ]\n",
    "                }\n",
    "            ]\n",
    "            text = processor.apply_chat_template(messages, add_generation_prompt=False)\n",
    "            texts.append(text.strip())\n",
    "\n",
    "        batch = processor(text=texts, return_tensors=\"pt\", padding=True)\n",
    "\n",
    "        #labels = batch[\"input_ids\"].clone()\n",
    "        #labels[labels == processor.tokenizer.pad_token_id] = self.image_token_id\n",
    "        #batch[\"labels\"] = labels\n",
    "\n",
    "        return batch\n",
    "\n",
    "data_collator = MyDataCollator(processor)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8f2e613-5695-4558-9a13-66158d82bed9",
   "metadata": {
    "id": "vsq4TtIJAMtH"
   },
   "source": [
    "### Setup training parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3cda658-05f6-4078-8d71-2d1c0352ecfa",
   "metadata": {
    "executionInfo": {
     "elapsed": 1008,
     "status": "ok",
     "timestamp": 1730998601172,
     "user": {
      "displayName": "Ronan Le Meillat",
      "userId": "09161391957806824350"
     },
     "user_tz": -60
    },
    "id": "Q_WKQFfoAMtH"
   },
   "outputs": [],
   "source": [
    "from transformers import TrainingArguments, Trainer\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir = output_dir,\n",
    "    overwrite_output_dir = False,\n",
    "    auto_find_batch_size = True,\n",
    "    learning_rate = 2e-4,\n",
    "    fp16 = True,\n",
    "    per_device_train_batch_size = 2,\n",
    "    per_device_eval_batch_size = 2,\n",
    "    gradient_accumulation_steps = 8,\n",
    "    dataloader_pin_memory = False,\n",
    "    save_total_limit = 3,\n",
    "    eval_strategy = \"steps\",\n",
    "    save_strategy = \"steps\",\n",
    "    eval_steps = 100,\n",
    "    save_steps = 10, # checkpoint each 10 steps\n",
    "    resume_from_checkpoint = True,\n",
    "    logging_steps = 5,\n",
    "    remove_unused_columns = False,\n",
    "    push_to_hub = True,\n",
    "    label_names = [\"labels\"],\n",
    "    load_best_model_at_end = False,\n",
    "    report_to = \"none\",\n",
    "    optim = \"paged_adamw_8bit\",\n",
    "    max_steps = 10, # remove this for training\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6569265-5941-4482-84e2-faf1b61b685c",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 426,
     "status": "ok",
     "timestamp": 1730998605441,
     "user": {
      "displayName": "Ronan Le Meillat",
      "userId": "09161391957806824350"
     },
     "user_tz": -60
    },
    "id": "vSIo17mgAMtH",
    "outputId": "3bebd35a-ed7f-49ee-e1bc-91594e8dcd24"
   },
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model = model,\n",
    "    args = training_args,\n",
    "    data_collator = data_collator,\n",
    "    train_dataset = train_dataset,\n",
    "    eval_dataset = eval_dataset,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "916ac153-206b-488a-b783-3ad0c4ba21b6",
   "metadata": {
    "id": "pmlwDsOpAMtI"
   },
   "source": [
    "### Start (or restart) Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb72a570-97e8-440e-b79f-640d8898e37c",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "WQA84KnTAMtI",
    "outputId": "ebb15160-f56e-4899-e608-b0d5fd0ba117"
   },
   "outputs": [],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b109f2b9-f3cb-4732-8318-b74ed9e5aa25",
   "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.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}