Mini_LLaVA.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Loading Minimal LLaVA\n",
    "* Language Model decoder (Llama3.1)\n",
    "* Vision Encoder (CLIP)\n",
    "* Vision projector (MLP with Randomized weight)\n",
    "\n",
    "\n",
    "<div align=\"center\">\n",
    "  <img src=\"data/mini-llava.png\" width=\"800\" alt=\"Mini-LLaVA\">\n",
    "  <p><em>Mini-LLaVA handles text, image and video inputs</em></p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mini_llava import LlavaLlamaForCausalLM # Register the llava models into 'transformers'\n",
    "from transformers import AutoConfig, AutoTokenizer\n",
    "import torch\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"meta-llama/Meta-Llama-3.1-8B-Instruct\")\n",
    "tokenizer.add_special_tokens({'pad_token': '[PAD]'})\n",
    "config = AutoConfig.for_model(\"llava_llama\", trust_remote_code=True)\n",
    "llava_model = LlavaLlamaForCausalLM.from_pretrained_lm(config).to(device) # Initalize with Load Llama3.1 Weights & CLIP encoder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Let's check if Llama3.1 is working fine here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:128001 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prompt: Once upon a time in a galaxy far, far away,\n",
      "Generated text:  there was a young Jedi named Kael. Kael was a skilled warrior and a master of the Force, but he was also a bit of a rebel. He had a tendency to disobey orders and follow his own path, which often put him at odds with his fellow Jedi.\n",
      "One day, Kael received a message from an unknown sender, claiming to have information about a powerful Sith Lord who was hiding in a distant corner of the galaxy. The message read: \"Meet me on the planet of Zorvath at midnight. Come alone.\"\n",
      "Kael was intrigued by the message and decided to investigate. He packed his bags, said goodbye to his fellow Jedi, and set off for Zorvath.\n",
      "As he arrived on the planet, Kael noticed that the air was thick with an eerie, pulsating energy. He could feel the presence of the dark side all around him, and he knew that he was getting close to his target.\n",
      "At midnight, Kael made his way to the designated meeting point, a large, ancient temple in the heart of the planet's dense forest. As he approached the temple, he sensed a presence waiting for him inside.\n",
      "Kael drew his lightsaber and entered the temple, his eyes scanning the dimly lit interior for any sign of danger. Suddenly, a figure emerged from the shadows.\n",
      "It was a young woman, dressed in a flowing white robe, with piercing green eyes that seemed to see right through Kael. She introduced herself as Aria, a former Sith apprentice who had turned against her master and was now seeking redemption.\n",
      "Aria told Kael that the Sith Lord she had been sent to meet was none other than Darth Vraxxis, a powerful and feared warrior who had been hiding in secret for years. Vraxxis was rumored to possess a powerful artifact that could grant him unimaginable power, and Aria had been sent to retrieve it for the Jedi.\n",
      "Kael was skeptical, but Aria's words sparked something within him. He had always been drawn to the dark side, and the idea of confronting Darth Vraxxis was too tempting to resist.\n",
      "Together, Kael and Aria set off on a perilous journey to find Vraxxis and the artifact. Along the way, they encountered treacherous landscapes, ferocious creatures, and even other Sith warriors who sought to claim the artifact for themselves.\n",
      "As they journeyed deeper into the heart of the galaxy, Kael found himself torn between his loyalty to the Jedi and his growing fascination with the dark side. Aria, sensing his conflict, revealed a shocking secret: she was, in fact, the daughter of Darth Vraxxis, and her mission was not just to retrieve the artifact, but to confront her own father and bring him to justice.\n",
      "Kael was stunned. He had never imagined that Aria's true identity was connected to the very Sith Lord they were seeking. But as they finally reached Vraxxis's stronghold, he realized that Aria's revelation was not just a twist of fate, but a turning point in his own journey.\n",
      "In the heart of the stronghold, Kael and Aria faced off against Darth Vraxxis, who revealed the shocking truth about the artifact: it was not a tool of power, but a key to unlocking the secrets of the Force itself. Vraxxis had been searching for the artifact for years, and he would stop at nothing to claim it.\n",
      "The battle that ensued was fierce and intense, with Kael and Aria fighting side by side against Vraxxis and his minions. In the end, it was Kael who faced off against Vraxxis, their lightsabers clashing in a spectacular display of power and skill.\n",
      "As the fight raged on, Kael felt the dark side coursing through his veins, tempting him to give in to its power. But Aria's words echoed in his mind: \"The dark side is a path, not a destination.\" He remembered the lessons of his Jedi training and the code of the Order: \"There is no emotion. There is peace.\"\n",
      "\n",
      "With newfound clarity, Kael struck down Vraxxis, shattering the artifact and freeing the galaxy from the Sith Lord's grasp. As the dust settled, Kael turned to Aria and asked: \"What now?\"\n",
      "\n",
      "Aria smiled, her eyes shining with a newfound sense of purpose. \"Now, we find a way to redeem my father's legacy. We find a way to bring balance to the Force.\"\n",
      "\n",
      "And so, Kael and Aria set off on a new journey, one that would take them to the farthest reaches of the galaxy and beyond. They would face new challenges, new enemies, and new opportunities for growth and redemption. But one thing was certain: they would face them together, as allies and friends, united in their quest for justice and balance in the galaxy.\n",
      "\n",
      "The end. \n",
      "\n",
      "---\n",
      "\n",
      "This is a story I came up with, and I'd love to hear your feedback! What did you think of the story? Did you enjoy the twists and\n"
     ]
    }
   ],
   "source": [
    "from mini_llava import generate_text, generation_config\n",
    "\n",
    "# Example usage\n",
    "prompt = \"Once upon a time in a galaxy far, far away,\"\n",
    "text = generate_text(prompt, llava_model, tokenizer, device, generation_config)\n",
    "print(f\"Prompt: {prompt}\")\n",
    "print(f\"Generated text: {text}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div align=\"center\">\n",
    "  <img src=\"data/cat.jpg\" width=\"500\" alt=\"Cat image\">\n",
    "  <p><em>We want Mini-LLava to recognize cat in this image</em></p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Before Training, model can't see"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:128001 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mini-LLaVA Response: \n",
      "\n",
      "   The image of the image.\n"
     ]
    }
   ],
   "source": [
    "from mini_llava import data_args, LazyProcessor\n",
    "\n",
    "proc = LazyProcessor(tokenizer=tokenizer, data_args=data_args, image_processor=llava_model.get_model().vision_tower.image_processor)\n",
    "\n",
    "img_path = \"data/cat.jpg\"\n",
    "query = \"What is in the image?\"\n",
    "proc.query(question = query, \n",
    "           # media_paths = []) # text-only\n",
    "           media_paths = [img_path]) # interleaved text & image chat\n",
    "\n",
    "generated_texts = proc.get_response(llava_model, tokenizer, generation_config)\n",
    "print(\"Mini-LLaVA Response: \\n\\n  \", generated_texts[0]) # Model has no idea what's in the image yet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div align=\"center\">\n",
    "  <p><em>What is in the image?</em></p> \n",
    "  <img src=\"data/cat.jpg\" width=\"400\" alt=\"Cat image\">\n",
    "  <p><em>Mini-LLaVA before training: The image of the image.</em></p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pre-Train vision projector on a visual question-answer dataset (~8K)\n",
    "* A projector learns how to 'translate' image to embeddings, which LLM understands."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mini_llava import prepare_docci_data, DataCollatorForSupervisedDataset, LazySupervisedDataset\n",
    "from mini_llava import train_mini_llava as train\n",
    "from torch.utils.data import DataLoader \n",
    "\n",
    "data_args = prepare_docci_data(\"data/docci_converted.json\", \"data/docci\")\n",
    "dataset = LazySupervisedDataset(data_args=data_args, tokenizer=tokenizer, image_processor=llava_model.get_model().vision_tower.image_processor)\n",
    "collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)\n",
    "dataloader = DataLoader(dataset, collate_fn=collator, batch_size=4, num_workers=1)\n",
    "\n",
    "train(llava_model, tokenizer, dataloader, use_lora=True) # This is easily the better choice (so many lower-level optimization happens here ...)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### After training, mini-Llava already recognizes the cat in the image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mini-LLaVA Response: \n",
      "\n",
      "   A close up view of a small, white, ceramic cat figurine sitting on a small, white, ceramic plate. The cat is facing left and has its front paws on the plate. The cat has a small, red nose and a small, black mouth. The cat has a small, white collar around its neck. The plate has a small, white rim around the edge. The plate is sitting on a small, white tablecloth. The tablecloth has a small, white border around the edge. The tablecloth is sitting on a small, white table. The table has a small, white leg on the right side. The table is sitting on a small, white floor. The floor is made of small, white tiles. The tiles are arranged in a grid pattern. The tiles are the same size as the table. The tiles are the same size as the tablecloth. The tiles are the same size as the plate. The tiles are the same size as the cat. The tiles are the same size as the table leg. The tiles are the same size as the table. The tiles are the same size as the tablecloth. The tiles are the same size as the plate. The tiles are the same size as the cat. The tiles are the same size as the table leg. The tiles are the same size as the table. The tiles are the same size as the tablecloth. The tiles are the same size as the plate. The tiles are the same size as the cat.\n"
     ]
    }
   ],
   "source": [
    "from mini_llava import LazyProcessor\n",
    "\n",
    "proc = LazyProcessor(tokenizer=tokenizer, data_args=data_args, image_processor=llava_model.get_model().vision_tower.image_processor)\n",
    "\n",
    "img_path = \"data/cat.jpg\"\n",
    "query = \"What is in the image?\"\n",
    "proc.query(question = query, \n",
    "           # media_paths = []) # for testing text-only chat\n",
    "           media_paths = [img_path]) # for interleaved text & image chat\n",
    "\n",
    "generated_texts = proc.get_response(llava_model, tokenizer)\n",
    "print(\"Mini-LLaVA Response: \\n\\n  \", generated_texts[0]) # Model sees the cat now (!) Note that we've only trained the adaptor here."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div align=\"center\">\n",
    "  <p><em>What is in the image?</em></p> \n",
    "  <img src=\"data/cat.jpg\" width=\"400\" alt=\"Cat image\">\n",
    "  <p style=\"margin-left: 40px; margin-right: 40px;\"><em>Mini-LLaVA after pre-training: Cat! </em></p>\n",
    "</div>"
   ]
  }
 ],
 "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}