{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "084f56a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'\n",
    "import sys\n",
    "import torch\n",
    "from PIL import Image\n",
    "import requests\n",
    "import json\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import scipy.io as sio\n",
    "import h5py\n",
    "from tqdm import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "import utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "837458e5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/py312/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "The image processor of type `Qwen2VLImageProcessor` is now loaded as a fast processor by default, even if the model checkpoint was saved with a slow processor. This is a breaking change and may produce slightly different outputs. To continue using the slow processor, instantiate this class with `use_fast=False`. Note that this behavior will be extended to all models in a future release.\n",
      "You have video processor config saved in `preprocessor.json` file which is deprecated. Video processor configs should be saved in their own `video_preprocessor.json` file. You can rename the file or load and save the processor back which renames it automatically. Loading from `preprocessor.json` will be removed in v5.0.\n",
      "Loading checkpoint shards: 100%|██████████| 18/18 [04:58<00:00, 16.58s/it]\n"
     ]
    }
   ],
   "source": [
    "from transformers import Qwen2_5_VLForConditionalGeneration, AutoTokenizer, AutoProcessor, BitsAndBytesConfig\n",
    "\n",
    "processor = AutoProcessor.from_pretrained(\"/opt/data/private/huggingface/models--Qwen--Qwen2.5-VL-32B-Instruct\")\n",
    "\n",
    "quant_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,  \n",
    "    bnb_4bit_quant_type=\"nf4\",   \n",
    "    bnb_4bit_compute_dtype=\"float16\" \n",
    ")\n",
    "model = Qwen2_5_VLForConditionalGeneration.from_pretrained(\n",
    "    \"/opt/data/private/huggingface/models--Qwen--Qwen2.5-VL-32B-Instruct\", device_map=\"auto\", quantization_config=quant_config\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "be5d46fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "nsddir = '/opt/data/private/dataset/nsd'\n",
    "stim_info = pd.read_csv(os.path.join(nsddir,'nsddata','experiments','nsd', 'nsd_stim_info_merged.csv'), index_col=0)\n",
    "cocoId = stim_info['cocoId']\n",
    "\n",
    "with open(os.path.join(nsddir, 'nsddata_stimuli/stimuli/nsd/annotations', 'nsd_captions.json'), 'r') as f:\n",
    "    captions = json.load(f)\n",
    "\n",
    "stimulus = h5py.File(os.path.join(nsddir,'nsddata_stimuli','stimuli','nsd','coco_images_224_float16.hdf5'), 'r')['images']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "id": "219c7f64",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['A cat sits on the edge of a toilet.',\n",
       " 'A cat sitting on the edge of the toilet looking toward the open bathroom door.',\n",
       " 'a black and gray cat is sitting on a toilet',\n",
       " 'A cat laying on a toilet seat with its paws inside .',\n",
       " 'A cat is laying on a toilet seat']"
      ]
     },
     "execution_count": 122,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 234280 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "imgidx = 1\n",
    "plt.imshow(utils.np_to_Image(stimulus[imgidx]))\n",
    "plt.axis('off')\n",
    "plt.grid(False)\n",
    "captions[str(cocoId[imgidx])]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 123,
   "id": "b96fbf34",
   "metadata": {},
   "outputs": [],
   "source": [
    "def gen_prompt(caption):\n",
    "    \"\"\"Generate a prompt for the model.\"\"\"\n",
    "    messages = [\n",
    "        {\n",
    "            \"role\": \"system\",\n",
    "            \"content\": \"You are an AI visual assistant that can analyze objects in the image. Currently, you receive an image and some sentences each describing the image you are observing. \\n Please describe objects and relevance, concepts, background, color and scene of the image in a detailed manner but without decoration and embellishment. \\n Always answer as if you are directly looking at the image. \\n Describe the image content clearly and concisely and retain the meanings of each objects, relevance, concepts, background, color and scene in the image. \\n  Describe the image directly from the beginning. Do not with 'The image shows' or 'The image depicts'. Don't summarize or overall. \\n Keep your answer less than 77 characters and words.\"\n",
    "        },\n",
    "        {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": [\n",
    "            {\n",
    "                \"type\": \"text\", \n",
    "                \"text\": caption\n",
    "            },\n",
    "            {\"type\": \"image\"},\n",
    "            ],\n",
    "        }\n",
    "    ]\n",
    "    prompt = processor.apply_chat_template(messages, add_generation_prompt=True)\n",
    "    return prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "id": "21ec48bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "for n in range(1, len(stimulus)):\n",
    "    images = [utils.np_to_Image(stimulus[n])]\n",
    "    prompts = [gen_prompt(captions[str(cocoId[n])])]\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "id": "48da4bf0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/py312/lib/python3.12/site-packages/transformers/generation/utils.py:2505: UserWarning: You are calling .generate() with the `input_ids` being on a device type different than your model's device. `input_ids` is on cpu, whereas the model is on cuda. You may experience unexpected behaviors or slower generation. Please make sure that you have put `input_ids` to the correct device by calling for example input_ids = input_ids.to('cuda') before running `.generate()`.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "inputs = processor(\n",
    "    text=prompts,\n",
    "    images=images,\n",
    "    padding=True,\n",
    "    return_tensors=\"pt\",\n",
    ")\n",
    "generated_ids = model.generate(**inputs, max_new_tokens=77, temperature=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 127,
   "id": "e59f8837",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"A black and gray cat sits on the edge of a white toilet, facing left toward an open bathroom door. The scene includes a yellow mat on the floor and a shower curtain in the background. The cat's tail hangs down.\""
      ]
     },
     "execution_count": 127,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "generated_ids_trimmed = [\n",
    "    out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)\n",
    "]\n",
    "output_text = processor.batch_decode(\n",
    "    generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False\n",
    ")\n",
    "output_text[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5ae92ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = {}\n",
    "data.update({str(cocoId[n]): output_text})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "id": "202045a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(nsddir, 'nsddata_stimuli/stimuli/nsd/annotations', 'qwen2_5_vl_captions.json'), 'r') as f:\n",
    "    old_data = json.load(f)\n",
    "    old_data.update(data)\n",
    "with open(os.path.join(nsddir, 'nsddata_stimuli/stimuli/nsd/annotations', 'qwen2_5_vl_captions.json'), 'w', encoding=\"utf-8\") as f:\n",
    "    json.dump(old_data, f, ensure_ascii=False, indent=4)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py312",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
