finetune_image_captioning_peft.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UNCNPVi8iAgw"
      },
      "source": [
        "# IDEFICS: A Flamingo-based model, trained at scale for the community\n",
        "# Finetuning Demo Notebook:\n",
        "\n",
        "\n",
        "Credit: [Flamingo blog](https://www.deepmind.com/blog/tackling-multiple-tasks-with-a-single-visual-language-model)\n",
        "\n",
        "This google colab notebook shows how to run predictions with the 4-bit quantized 🤗 [Idefics-9B model](https://huggingface.co/HuggingFaceM4/idefics-9b) and finetune it on a specific dataset.\n",
        "\n",
        "[IDEFICS](https://huggingface.co/HuggingFaceM4/idefics-80b) is a multi-modal model based on the [Flamingo](https://arxiv.org/abs/2204.14198) architecture. It can take images and texts as input and return text outputs but it does not support image generation. \\\\\n",
        "IDEFICS is built on top of two unimodal open-access pre-trained models to connect the two modalities. Newly initialized parameters in the form of Transformer blocks bridge the gap between the vision encoder and the language model. The model is trained on a mixture of image/text pairs and unstrucutred multimodal web documents. \\\\\n",
        "The [finetuned versions](https://huggingface.co/HuggingFaceM4/idefics-80b-instruct) of IDEFICS behave like LLM chatbots while also understanding visual input. \\\\\n",
        "You can play with the [demo here](https://huggingface.co/spaces/HuggingFaceM4/idefics_playground)\n",
        "\n",
        "The code for this notebook was contributed to by *Léo Tronchon, Younes Belkada, and Stas Bekman*, the IDEFICS model has been contributed to by: *Lucile Saulnier, Léo Tronchon, Hugo Laurençon, Stas Bekman, Amanpreet Singh, Siddharth Karamcheti, and Victor Sanh*"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7m9zw1wcCC8e"
      },
      "source": [
        "# Install and import necessary libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "prXRsUiXCII9"
      },
      "outputs": [],
      "source": [
        "!pip install -q datasets\n",
        "!pip install -q git+https://github.com/huggingface/transformers.git\n",
        "!pip install -q bitsandbytes sentencepiece accelerate loralib\n",
        "!pip install -q -U git+https://github.com/huggingface/peft.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MxoHmx-HfAgf"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from datasets import load_dataset\n",
        "from peft import LoraConfig, get_peft_model\n",
        "from PIL import Image\n",
        "from transformers import IdeficsForVisionText2Text, AutoProcessor, Trainer, TrainingArguments, BitsAndBytesConfig\n",
        "import torchvision.transforms as transforms"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DP_ilre6jI6l"
      },
      "source": [
        "# Load quantized model\n",
        "First get the quantized version of the model. This will allow us to use the 9B version of Idefics with a single 16GB gpu\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U8qLbD1Iocak"
      },
      "outputs": [],
      "source": [
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IRiT0q0Ck-3Y"
      },
      "outputs": [],
      "source": [
        "# checkpoint = \"HuggingFaceM4/tiny-random-idefics\"\n",
        "checkpoint = \"HuggingFaceM4/idefics-9b\"\n",
        "\n",
        "# Here we skip some special modules that can't be quantized properly\n",
        "bnb_config = BitsAndBytesConfig(\n",
        "    load_in_4bit=True,\n",
        "    bnb_4bit_use_double_quant=True,\n",
        "    bnb_4bit_quant_type=\"nf4\",\n",
        "    bnb_4bit_compute_dtype=torch.float16,\n",
        "    llm_int8_skip_modules=[\"lm_head\", \"embed_tokens\"],\n",
        ")\n",
        "\n",
        "processor = AutoProcessor.from_pretrained(checkpoint, use_auth_token=False)\n",
        "# Simply take-off the quantization_config arg if you want to load the original model\n",
        "model = IdeficsForVisionText2Text.from_pretrained(checkpoint, quantization_config=bnb_config, device_map=\"auto\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PloVmAaM75kJ"
      },
      "source": [
        "If you print the model, you will see that all `nn.Linear` layers are in fact replaced by `bnb.nn.Linear4bit` layers."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4gaDzRK174Ur"
      },
      "outputs": [],
      "source": [
        "print(device)\n",
        "print(model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8EIvXR6fPG6d"
      },
      "source": [
        "# Inference\n",
        "Let's make a simple method to test the model's inference"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J5MSZ3xdPF4f"
      },
      "outputs": [],
      "source": [
        "def check_inference(model, processor, prompts, max_new_tokens=50):\n",
        "    tokenizer = processor.tokenizer\n",
        "    bad_words = [\"<image>\", \"<fake_token_around_image>\"]\n",
        "    if len(bad_words) > 0:\n",
        "        bad_words_ids = tokenizer(bad_words, add_special_tokens=False).input_ids\n",
        "\n",
        "    eos_token = \"</s>\"\n",
        "    eos_token_id = tokenizer.convert_tokens_to_ids(eos_token)\n",
        "\n",
        "    inputs = processor(prompts, return_tensors=\"pt\").to(device)\n",
        "    generated_ids = model.generate(**inputs, eos_token_id=[eos_token_id], bad_words_ids=bad_words_ids, max_new_tokens=max_new_tokens, early_stopping=True)\n",
        "    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]\n",
        "    print(generated_text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RYA2HKGC0n9d"
      },
      "source": [
        "\n",
        "Let's run prediction with the quantized model for the image below which pictures two kittens. \\\\\n",
        "<img src=\"https://hips.hearstapps.com/hmg-prod/images/cute-photos-of-cats-in-grass-1593184777.jpg\" width=\"400\"/>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6I_iDtQN03jE"
      },
      "outputs": [],
      "source": [
        "url = \"https://hips.hearstapps.com/hmg-prod/images/cute-photos-of-cats-in-grass-1593184777.jpg\"\n",
        "prompts = [\n",
        "    # \"Instruction: provide an answer to the question. Use the image to answer.\\n\",\n",
        "    url,\n",
        "    \"Question: What's on the picture? Answer:\",\n",
        "]\n",
        "check_inference(model, processor, prompts, max_new_tokens=5)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DLiwPnGBxiJf"
      },
      "source": [
        "Now let's see how the model fares on pokemon knowledge before we try to finetune it further. \\\\\n",
        "<img src=\"https://images.pokemontcg.io/pop6/2_hires.png\" width=\"194\"/>\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lDVDUE1ew7tZ"
      },
      "outputs": [],
      "source": [
        "# check generation before finetuning\n",
        "\n",
        "url = \"https://images.pokemontcg.io/pop6/2_hires.png\"\n",
        "prompts = [\n",
        "    url,\n",
        "    \"Question: What's on the picture? Answer:\",\n",
        "]\n",
        "check_inference(model, processor, prompts, max_new_tokens=100)\n",
        "# It looks like the model is already aware of pokemon - but it could be more specific, and less repetitive"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ydBhQT6SQiWy"
      },
      "source": [
        "# Finetuning dataset\n",
        "Prepare the dataset that will be used for finetuning\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PWsgZRnJXmJC"
      },
      "outputs": [],
      "source": [
        "import requests\n",
        "def is_image_url(url):\n",
        "    try:\n",
        "        response = requests.head(url, allow_redirects=True)\n",
        "        content_type = response.headers.get('Content-Type')\n",
        "        if content_type and 'image' in content_type.lower():\n",
        "            return True\n",
        "        else:\n",
        "            return False\n",
        "    except requests.RequestException as e:\n",
        "        print(f\"Error checking URL: {e}\")\n",
        "        return False\n",
        "\n",
        "def convert_to_rgb(image):\n",
        "    # `image.convert(\"RGB\")` would only work for .jpg images, as it creates a wrong background\n",
        "    # for transparent images. The call to `alpha_composite` handles this case\n",
        "    if image.mode == \"RGB\":\n",
        "        return image\n",
        "\n",
        "    image_rgba = image.convert(\"RGBA\")\n",
        "    background = Image.new(\"RGBA\", image_rgba.size, (255, 255, 255))\n",
        "    alpha_composite = Image.alpha_composite(background, image_rgba)\n",
        "    alpha_composite = alpha_composite.convert(\"RGB\")\n",
        "    return alpha_composite\n",
        "\n",
        "def ds_transforms(example_batch):\n",
        "    image_size = processor.image_processor.image_size\n",
        "    image_mean = processor.image_processor.image_mean\n",
        "    image_std = processor.image_processor.image_std\n",
        "\n",
        "    image_transform = transforms.Compose([\n",
        "        convert_to_rgb,\n",
        "        transforms.RandomResizedCrop((image_size, image_size), scale=(0.9, 1.0), interpolation=transforms.InterpolationMode.BICUBIC),\n",
        "        transforms.ToTensor(),\n",
        "        transforms.Normalize(mean=image_mean, std=image_std),\n",
        "    ])\n",
        "\n",
        "    prompts = []\n",
        "    for i in range(len(example_batch['entity_name'])):\n",
        "        # We split the captions to avoid having very long examples, which would require more GPU ram during training\n",
        "        # caption = example_batch['caption'][i].split(\".\")[0]\n",
        "        if is_image_url(example_batch['image_link'][i]):\n",
        "          prompts.append(\n",
        "              [\n",
        "                  example_batch['image_link'][i],\n",
        "                  # f\"Question: What's the {example_batch['entity_name'][i]} of the . {caption}</s>\",\n",
        "                  f\"Question: What's numerical {example_batch['entity_name'][i]} shown in the picture of the object, don't modify? Answer: {example_batch['entity_value'][i]}\",\n",
        "              ],\n",
        "          )\n",
        "        else:\n",
        "          print(\"image doesn't exist here: \", example_batch['image_link'][i])\n",
        "\n",
        "    inputs = processor(prompts, transform=image_transform, return_tensors=\"pt\").to(device)\n",
        "\n",
        "    inputs[\"labels\"] = inputs[\"input_ids\"]\n",
        "\n",
        "    return inputs\n",
        "\n",
        "\n",
        "# load and prepare dataset\n",
        "# Load your CSV file into a DatasetDict\n",
        "ds = load_dataset('csv', data_files={'train': '/content/train.csv'})\n",
        "# ds = load_dataset(\"TheFusion21/PokemonCards\")\n",
        "ds = ds[\"train\"].train_test_split(test_size=0.002)\n",
        "train_ds = ds[\"train\"]\n",
        "eval_ds = ds[\"test\"]\n",
        "train_ds.set_transform(ds_transforms)\n",
        "eval_ds.set_transform(ds_transforms)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U2JwmYH8W9-t"
      },
      "outputs": [],
      "source": [
        "print(type(ds))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ibx9_fLUdQ8F"
      },
      "outputs": [],
      "source": [
        "train_ds"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Kui4EkCmOQzd"
      },
      "source": [
        "# LoRA\n",
        "After specifying the low-rank adapters (LoRA) config, we load the PeftModel using the get_peft_model utility function"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Zllx_apSd_ac"
      },
      "outputs": [],
      "source": [
        "checkpoint"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jKa5oTorp_A-"
      },
      "outputs": [],
      "source": [
        "model_name = checkpoint.split(\"/\")[1]\n",
        "config = LoraConfig(\n",
        "    r=16,\n",
        "    lora_alpha=32,\n",
        "    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\"],\n",
        "    lora_dropout=0.05,\n",
        "    bias=\"none\",\n",
        ")\n",
        "model = get_peft_model(model, config)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ShuZJ5K2pYoL"
      },
      "outputs": [],
      "source": [
        "model.print_trainable_parameters()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0Ok1sOZKQ29s"
      },
      "source": [
        "# Training\n",
        "Finally, using the Hugging Face Trainer, we can finetune the model!\n",
        "\n",
        "For the sake of the demo, we have set the max_steps at 40. That's about 0.05 epoch on this dataset, so feel free to tune further!\n",
        "\n",
        "It has been reported that fine-tuning in mixed precision fp16 can lead to overflows. As such, we recommend training in mixed precision bf16 when possible."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9cD3OuygpR5l"
      },
      "outputs": [],
      "source": [
        "training_args = TrainingArguments(\n",
        "    output_dir=f\"{model_name}-amazon-images\",\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",
        "    evaluation_strategy=\"steps\",\n",
        "    save_strategy=\"steps\",\n",
        "    save_steps=40,\n",
        "    eval_steps=20,\n",
        "    logging_steps=20,\n",
        "    max_steps=40,\n",
        "    remove_unused_columns=False,\n",
        "    push_to_hub=False,\n",
        "    label_names=[\"labels\"],\n",
        "    load_best_model_at_end=True,\n",
        "    report_to=None,\n",
        "    optim=\"paged_adamw_8bit\",\n",
        ")\n",
        "\n",
        "trainer = Trainer(\n",
        "    model=model,\n",
        "    args=training_args,\n",
        "    train_dataset=train_ds,\n",
        "    eval_dataset=eval_ds,\n",
        ")\n",
        "\n",
        "trainer.train()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oQLNiiLNpQYT"
      },
      "outputs": [],
      "source": [
        "# Save the model weights\n",
        "trainer.save_model(f\"/content/{model_name}-amazon-images-ml\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tkv3wBsjOxUd"
      },
      "source": [
        "## Inference from local model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "lXFkmuRgL59x"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoProcessor\n",
        "checkpoint = \"HuggingFaceM4/idefics-9b\"\n",
        "processor = AutoProcessor.from_pretrained(checkpoint, use_auth_token=False)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HSrooNAaJuby"
      },
      "outputs": [],
      "source": [
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "from transformers import IdeficsForVisionText2Text\n",
        "saved_model_path = f\"/content/idefics-9b-amazon-images-ml\"\n",
        "# Load the processor (tokenizer, image processor, etc.)\n",
        "# processor_ = AutoProcessor.from_pretrained(saved_model_path)\n",
        "\n",
        "# Here we skip some special modules that can't be quantized properly\n",
        "bnb_config = BitsAndBytesConfig(\n",
        "    load_in_4bit=True,\n",
        "    bnb_4bit_use_double_quant=True,\n",
        "    bnb_4bit_quant_type=\"nf4\",\n",
        "    bnb_4bit_compute_dtype=torch.float16,\n",
        "    llm_int8_skip_modules=[\"lm_head\", \"embed_tokens\"],\n",
        ")\n",
        "\n",
        "# Load the model\n",
        "model_ = IdeficsForVisionText2Text.from_pretrained(saved_model_path, quantization_config=bnb_config).to(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "# Set the model to evaluation mode\n",
        "model_.eval()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "C07a8ML7hDW_"
      },
      "outputs": [],
      "source": [
        "ppt = [\n",
        "    \"https://m.media-amazon.com/images/I/51xvxFlRXoL.jpg\",\n",
        "    # f\"Question: What's the {example_batch['entity_name'][i]} of the . {caption}</s>\",\n",
        "    \"Question: What's numerical height shown in the picture of the objct, don't modify?\",\n",
        "]\n",
        "check_inference(model, processor, ppt, max_new_tokens=100)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zgqonle8AdPs"
      },
      "source": [
        "# Push your new model to the hub!\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zsIJzQGhmKBe"
      },
      "outputs": [],
      "source": [
        "!git config --global credential.helper store\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KrnB4kFxAjIA"
      },
      "outputs": [],
      "source": [
        "# Insert your \"write\" token. You should find it in the settings of your HF profile\n",
        "!huggingface-cli login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KF3O8oCOTeig"
      },
      "source": [
        "## Upload Model to HuggingFace Hub\n",
        "## Download and inference"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VT6C0YrxwiAA"
      },
      "outputs": [],
      "source": [
        "from huggingface_hub import HfApi, HfFolder\n",
        "\n",
        "# Initialize the Hugging Face API\n",
        "api = HfApi()\n",
        "token = HfFolder.get_token()\n",
        "\n",
        "\n",
        "# Upload the model files\n",
        "api.upload_folder(\n",
        "    folder_path=\"/content/idefics-9b-amazon-images-ml\",  # Path to the folder containing model files\n",
        "    repo_id=\"Ashu777/idefics-9b-amazon-image-captioning\",\n",
        "    repo_type=\"model\",\n",
        "    token=token\n",
        ")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "omtywa5UEE-Q"
      },
      "outputs": [],
      "source": [
        "!pip install bitsandbytes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ChjgND8CTBmE"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from datasets import load_dataset\n",
        "from peft import LoraConfig, get_peft_model\n",
        "from PIL import Image\n",
        "from transformers import IdeficsForVisionText2Text, AutoProcessor, Trainer, TrainingArguments, BitsAndBytesConfig\n",
        "import torchvision.transforms as transforms"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SzYj42NVRKJE"
      },
      "outputs": [],
      "source": [
        "# checkpoint = \"HuggingFaceM4/tiny-random-idefics\"\n",
        "checkpoint = \"Ashu777/idefics-9b-amazon-image-captioning\"\n",
        "\n",
        "# Here we skip some special modules that can't be quantized properly\n",
        "bnb_config = BitsAndBytesConfig(\n",
        "    load_in_4bit=True,\n",
        "    bnb_4bit_use_double_quant=True,\n",
        "    bnb_4bit_quant_type=\"nf4\",\n",
        "    bnb_4bit_compute_dtype=torch.float16,\n",
        "    llm_int8_skip_modules=[\"lm_head\", \"embed_tokens\"],\n",
        ")\n",
        "\n",
        "# processor = AutoProcessor.from_pretrained(checkpoint, use_auth_token=False)\n",
        "# Simply take-off the quantization_config arg if you want to load the original model\n",
        "model = IdeficsForVisionText2Text.from_pretrained(checkpoint, quantization_config=bnb_config, device_map=\"auto\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Jb401iyZS2QV"
      },
      "outputs": [],
      "source": [
        "checkpoint = \"HuggingFaceM4/idefics-9b\"\n",
        "\n",
        "processor = AutoProcessor.from_pretrained(checkpoint, use_auth_token=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8vZsdButSuIR"
      },
      "outputs": [],
      "source": [
        "def check_inference(model, processor, prompts, max_new_tokens=50):\n",
        "    tokenizer = processor.tokenizer\n",
        "    bad_words = [\"<image>\", \"<fake_token_around_image>\"]\n",
        "    if len(bad_words) > 0:\n",
        "        bad_words_ids = tokenizer(bad_words, add_special_tokens=False).input_ids\n",
        "\n",
        "    eos_token = \"</s>\"\n",
        "    eos_token_id = tokenizer.convert_tokens_to_ids(eos_token)\n",
        "\n",
        "    inputs = processor(prompts, return_tensors=\"pt\").to(device)\n",
        "    generated_ids = model.generate(**inputs, eos_token_id=[eos_token_id], bad_words_ids=bad_words_ids, max_new_tokens=max_new_tokens, early_stopping=True)\n",
        "    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]\n",
        "    print(generated_text)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5HyAPTtRRKGG"
      },
      "outputs": [],
      "source": [
        "ppt = [\n",
        "    \"https://m.media-amazon.com/images/I/71lxjbXw9bL.jpg\",\n",
        "    # f\"Question: What's the {example_batch['entity_name'][i]} of the . {caption}</s>\",\n",
        "    \"Question: What's weight\",\n",
        "]\n",
        "check_inference(model, processor, ppt, max_new_tokens=100)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "25P9FD8_RKCz"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# inferencing at gradio"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "!pip install gradio"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sipFb0Yxyn0z"
      },
      "outputs": [],
      "source": [
        "import gradio as gr\n",
        "\n",
        "# Gradio interface function\n",
        "def run_inference(image_url, question):\n",
        "    ppt = [image_url, f\"Question: {question}\"]\n",
        "    return check_inference(model, processor, ppt, max_new_tokens=100)\n",
        "\n",
        "# Launch the Gradio app\n",
        "interface = gr.Interface(\n",
        "    fn=run_inference,\n",
        "    inputs=[gr.Textbox(label=\"Image URL\"), gr.Textbox(label=\"Question\")],\n",
        "    outputs=\"text\",\n",
        "    title=\"Vision-Text Model Inference\"\n",
        ")\n",
        "\n",
        "interface.launch()"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "private_outputs": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.4"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}