{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fb6407a0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[MultimodalReasoning] PyTorch is using primary device: cuda.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "796d05f5cb1c44608c9729c5ea1325c5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading weights:   0%|          | 0/105 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9d6a656936fc401f8d513395f05b29ba",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading weights:   0%|          | 0/103 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from transformers import AutoModel, AutoProcessor\n",
    "import models.visual as VisualModel\n",
    "import data.projection_dataset as ProjectionDataset\n",
    "import training.projector as ProjectorTraining\n",
    "import pipeline.Multimodal as Multimodal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c40755d0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[transformers] `torch_dtype` is deprecated! Use `dtype` instead!\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "af0c533f4b5d4167bf19b150a9698981",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading weights:   0%|          | 0/338 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "edcf914cea384efcbf6ae9ec354841bd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading weights:   0%|          | 0/792 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[VisualLayer] loaded from outputs/visual_projection_layer/visual_layer.pt\n"
     ]
    }
   ],
   "source": [
    "BaseMultimodal = Multimodal.MultimodalModel(\n",
    "    path=\"stable_models/reasoning_dpo_1\",\n",
    "    visual_path = 'outputs/visual_projection_layer',\n",
    "    is_trainable = False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f7c2dbd0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c1471804431a4bbf995ef034cafdcd66",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading weights:   0%|          | 0/792 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[VisualLayer] loaded from outputs/visual_projection_layer/visual_layer.pt\n"
     ]
    }
   ],
   "source": [
    "visual_model = VisualModel.VisualLayer(freeze_encoder=True)\n",
    "visual_model.load(\"outputs/visual_projection_layer\")\n",
    "trainer = ProjectorTraining.ProjectorTrainer(\n",
    "    language_model=BaseMultimodal.llm,\n",
    "    visual_model=visual_model,\n",
    "    lr=1e-4,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cf207ad",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I'm sorry, but I need more information to provide a clear description of what you're referring to as \"the image.\" Images can vary greatly in their content, such as photographs, paintings, drawings, or digital representations. Without knowing what specific image you have in mind, it's difficult for me to offer a precise response. If you could describe the image or provide some details about its characteristics, I'll do my best to help you with your request.\n"
     ]
    }
   ],
   "source": [
    "import PIL.Image as Image\n",
    "img = Image.open(\"dataset/living_room.png\").convert(\"RGB\")\n",
    "\n",
    "out = BaseMultimodal(\n",
    "    \"Describe the image given.\",\n",
    "    images=[img],\n",
    "    max_new_tokens=256,\n",
    "    retrieve=False\n",
    ")\n",
    "\n",
    "print(out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a2873b20",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Dataset] loaded 157712 from dataset/LLaVA-Instruct-150K/llava_instruct_150k.json\n",
      "[Dataset] total samples: 157712\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d35d2659941e4e7e8c18d497362ac9c9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 1/1:   0%|          | 0/36000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[VisualLayer] saved to outputs/visual_projection_layer\n",
      "[VisualLayer] saved to outputs/visual_projection_layer\n",
      "[VisualLayer] saved to outputs/visual_projection_layer\n",
      "[VisualLayer] saved to outputs/visual_projection_layer\n",
      "[VisualLayer] saved to outputs/visual_projection_layer\n",
      "[VisualLayer] saved to outputs/visual_projection_layer\n",
      "[VisualLayer] saved to outputs/visual_projection_layer\n",
      "[VisualLayer] saved to outputs/visual_projection_layer\n",
      "[VisualLayer] saved to outputs/visual_projection_layer\n",
      "[VisualLayer] saved to outputs/visual_projection_layer\n",
      "[VisualLayer] saved to outputs/visual_projection_layer\n",
      "[VisualLayer] saved to outputs/visual_projection_layer\n",
      "[VisualLayer] saved to outputs/visual_projection_layer\n",
      "[VisualLayer] saved to outputs/visual_projection_layer\n",
      "[VisualLayer] saved to outputs/visual_projection_layer\n",
      "[VisualLayer] saved to outputs/visual_projection_layer\n",
      "[VisualLayer] saved to outputs/visual_projection_layer\n",
      "[VisualLayer] saved to outputs/visual_projection_layer\n",
      "[VisualLayer] saved to outputs/visual_projection_layer\n",
      "[VisualLayer] saved to outputs/visual_projection_layer\n",
      "[VisualLayer] saved to outputs/visual_projection_layer\n",
      "[VisualLayer] saved to outputs/visual_projection_layer\n",
      "[VisualLayer] saved to outputs/visual_projection_layer\n",
      "[VisualLayer] saved to outputs/visual_projection_layer\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mKeyboardInterrupt\u001b[39m                         Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[5]\u001b[39m\u001b[32m, line 8\u001b[39m\n\u001b[32m      1\u001b[39m dataset = ProjectionDataset.InstructProjectorDataset(\n\u001b[32m      2\u001b[39m     projector_trainer=trainer,\n\u001b[32m      3\u001b[39m     json_paths=[\n\u001b[32m      4\u001b[39m         \u001b[33m\"\u001b[39m\u001b[33mdataset/LLaVA-Instruct-150K/llava_instruct_150k.json\u001b[39m\u001b[33m\"\u001b[39m,\n\u001b[32m      5\u001b[39m     ],\n\u001b[32m      6\u001b[39m )\n\u001b[32m----> \u001b[39m\u001b[32m8\u001b[39m \u001b[43mdataset\u001b[49m\u001b[43m.\u001b[49m\u001b[43mtrain\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m      9\u001b[39m \u001b[43m    \u001b[49m\u001b[43mepochs\u001b[49m\u001b[43m=\u001b[49m\u001b[32;43m1\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[32m     10\u001b[39m \u001b[43m    \u001b[49m\u001b[43mbatch_size\u001b[49m\u001b[43m=\u001b[49m\u001b[32;43m1\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[32m     11\u001b[39m \u001b[43m    \u001b[49m\u001b[43mmax_samples\u001b[49m\u001b[43m=\u001b[49m\u001b[32;43m36000\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[32m     12\u001b[39m \u001b[43m    \u001b[49m\u001b[43mshuffle\u001b[49m\u001b[43m=\u001b[49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[32m     13\u001b[39m \u001b[43m)\u001b[49m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/deep_learning/Projects/MultimodalReasoning/data/projection_dataset.py:200\u001b[39m, in \u001b[36mInstructProjectorDataset.train\u001b[39m\u001b[34m(self, epochs, max_samples, batch_size, shuffle, skip_missing, clear_every)\u001b[39m\n\u001b[32m    197\u001b[39m batch_answers.append(answer)\n\u001b[32m    199\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(batch_images) >= batch_size:\n\u001b[32m--> \u001b[39m\u001b[32m200\u001b[39m     loss = \u001b[43mflush_batch\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    201\u001b[39m     avg_loss = total_loss / \u001b[38;5;28mmax\u001b[39m(good_steps, \u001b[32m1\u001b[39m)\n\u001b[32m    203\u001b[39m     pbar.set_postfix({\n\u001b[32m    204\u001b[39m         \u001b[33m\"\u001b[39m\u001b[33mloss\u001b[39m\u001b[33m\"\u001b[39m: \u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mloss\u001b[38;5;132;01m:\u001b[39;00m\u001b[33m.4f\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m\"\u001b[39m,\n\u001b[32m    205\u001b[39m         \u001b[33m\"\u001b[39m\u001b[33mavg\u001b[39m\u001b[33m\"\u001b[39m: \u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mavg_loss\u001b[38;5;132;01m:\u001b[39;00m\u001b[33m.4f\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m\"\u001b[39m,\n\u001b[32m    206\u001b[39m         \u001b[33m\"\u001b[39m\u001b[33msteps\u001b[39m\u001b[33m\"\u001b[39m: good_steps,\n\u001b[32m    207\u001b[39m         \u001b[33m\"\u001b[39m\u001b[33mskipped\u001b[39m\u001b[33m\"\u001b[39m: skipped,\n\u001b[32m    208\u001b[39m     })\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/deep_learning/Projects/MultimodalReasoning/data/projection_dataset.py:164\u001b[39m, in \u001b[36mInstructProjectorDataset.train.<locals>.flush_batch\u001b[39m\u001b[34m()\u001b[39m\n\u001b[32m    161\u001b[39m good_steps += \u001b[32m1\u001b[39m\n\u001b[32m    162\u001b[39m global_step += \u001b[32m1\u001b[39m\n\u001b[32m--> \u001b[39m\u001b[32m164\u001b[39m \u001b[43mclear_batch\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    166\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m clear_every \u001b[38;5;129;01mand\u001b[39;00m global_step % clear_every == \u001b[32m0\u001b[39m:\n\u001b[32m    167\u001b[39m     cleanup_cuda()\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/deep_learning/Projects/MultimodalReasoning/data/projection_dataset.py:143\u001b[39m, in \u001b[36mInstructProjectorDataset.train.<locals>.clear_batch\u001b[39m\u001b[34m()\u001b[39m\n\u001b[32m    140\u001b[39m batch_questions = []\n\u001b[32m    141\u001b[39m batch_answers = []\n\u001b[32m--> \u001b[39m\u001b[32m143\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mclear_batch\u001b[39m():\n\u001b[32m    144\u001b[39m     batch_images.clear()\n\u001b[32m    145\u001b[39m     batch_questions.clear()\n",
      "\u001b[31mKeyboardInterrupt\u001b[39m: "
     ]
    }
   ],
   "source": [
    "\n",
    "dataset = ProjectionDataset.InstructProjectorDataset(\n",
    "    projector_trainer=trainer,\n",
    "    json_paths=[\n",
    "        \"dataset/LLaVA-Instruct-150K/llava_instruct_150k.json\",\n",
    "    ],\n",
    ")\n",
    "\n",
    "dataset.train(\n",
    "    epochs=1,\n",
    "    batch_size=1,\n",
    "    max_samples=36000,\n",
    "    shuffle=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1f2df73",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "TensorFlow 2.20 (WSL • GPU)",
   "language": "python",
   "name": "tf220-wsl-gpu"
  },
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}