{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import os\n",
    "import random\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.backends.cudnn as cudnn\n",
    "from tqdm import tqdm\n",
    "\n",
    "from torchvision import transforms\n",
    "from torchvision.transforms.functional import InterpolationMode\n",
    "from torchvision.utils import save_image\n",
    "\n",
    "from pope_loader import POPEDataSet\n",
    "from minigpt4.common.dist_utils import get_rank\n",
    "from minigpt4.models import load_preprocess\n",
    "\n",
    "from minigpt4.common.config import Config\n",
    "from minigpt4.common.dist_utils import get_rank\n",
    "from minigpt4.common.registry import registry\n",
    "\n",
    "# imports modules for registration\n",
    "from minigpt4.datasets.builders import *\n",
    "from minigpt4.models import *\n",
    "from minigpt4.processors import *\n",
    "from minigpt4.runners import *\n",
    "from minigpt4.tasks import *\n",
    "\n",
    "from PIL import Image\n",
    "from torchvision.utils import save_image\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "import seaborn\n",
    "import json\n",
    "\n",
    "\n",
    "MODEL_EVAL_CONFIG_PATH = {\n",
    "    \"minigpt4\": \"eval_configs/minigpt4_eval.yaml\",\n",
    "    \"instructblip\": \"eval_configs/instructblip_eval.yaml\",\n",
    "    \"lrv_instruct\": \"eval_configs/lrv_instruct_eval.yaml\",\n",
    "    \"shikra\": \"eval_configs/shikra_eval.yaml\",\n",
    "    \"llava-1.5\": \"eval_configs/llava-1.5_eval.yaml\",\n",
    "}\n",
    "\n",
    "POPE_PATH = {\n",
    "    \"random\": \"coco_pope/coco_pope_random.json\",\n",
    "    \"popular\": \"coco_pope/coco_pope_popular.json\",\n",
    "    \"adversarial\": \"coco_pope/coco_pope_adversarial.json\",\n",
    "}\n",
    "\n",
    "INSTRUCTION_TEMPLATE = {\n",
    "    \"minigpt4\": \"###Human: <Img><ImageHere></Img> <question> ###Assistant:\",\n",
    "    \"instructblip\": \"<ImageHere><question>\",\n",
    "    \"lrv_instruct\": \"###Human: <Img><ImageHere></Img> <question> ###Assistant:\",\n",
    "    \"shikra\": \"USER: <im_start><ImageHere><im_end> <question> ASSISTANT:\",\n",
    "    \"llava-1.5\": \"USER: <ImageHere> <question> ASSISTANT:\"\n",
    "}\n",
    "\n",
    "\n",
    "def setup_seeds(config):\n",
    "    seed = config.run_cfg.seed + get_rank()\n",
    "\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    cudnn.benchmark = False\n",
    "    cudnn.deterministic = True\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "parser = argparse.ArgumentParser(description=\"POPE-Adv evaluation on LVLMs.\")\n",
    "parser.add_argument(\"--model\", type=str, help=\"model\")\n",
    "parser.add_argument(\"--gpu-id\", type=int, help=\"specify the gpu to load the model.\")\n",
    "parser.add_argument(\n",
    "    \"--options\",\n",
    "    nargs=\"+\",\n",
    "    help=\"override some settings in the used config, the key-value pair \"\n",
    "    \"in xxx=yyy format will be merged into config file (deprecate), \"\n",
    "    \"change to --cfg-options instead.\",\n",
    ")\n",
    "parser.add_argument(\"--data_path\", type=str, default=\"/mnt/petrelfs/share_data/wangjiaqi/mllm-data-alg/COCO_2014/ori/val2014/val2014/\", help=\"data path\")\n",
    "parser.add_argument(\"--batch_size\", type=int, help=\"batch size\")\n",
    "parser.add_argument(\"--num_workers\", type=int, default=2, help=\"num workers\")\n",
    "args = parser.parse_known_args()[0]\n",
    "\n",
    "\n",
    "args.model = \"llava-1.5\"\n",
    "# args.model = \"instructblip\"\n",
    "# args.model = \"minigpt4\"\n",
    "# args.model = \"shikra\"\n",
    "args.gpu_id = \"0\"\n",
    "args.batch_size = 1\n",
    "\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = str(args.gpu_id)\n",
    "args.cfg_path = MODEL_EVAL_CONFIG_PATH[args.model]\n",
    "cfg = Config(args)\n",
    "setup_seeds(cfg)\n",
    "device = torch.device(\"cuda\") if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "# ========================================\n",
    "#             Model Initialization\n",
    "# ========================================\n",
    "print('Initializing Model')\n",
    "\n",
    "model_config = cfg.model_cfg\n",
    "model_config.device_8bit = args.gpu_id\n",
    "model_cls = registry.get_model_class(model_config.arch)\n",
    "model = model_cls.from_config(model_config).to(device)\n",
    "model.eval()\n",
    "processor_cfg = cfg.get_config().preprocess\n",
    "processor_cfg.vis_processor.eval.do_normalize = False\n",
    "vis_processors, txt_processors = load_preprocess(processor_cfg)\n",
    "print(vis_processors[\"eval\"].transform)\n",
    "print(\"Done!\")\n",
    "\n",
    "mean = (0.48145466, 0.4578275, 0.40821073)\n",
    "std = (0.26862954, 0.26130258, 0.27577711)\n",
    "norm = transforms.Normalize(mean, std)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_path = \"XXX.jpg\"\n",
    "raw_image = Image.open(image_path)\n",
    "plt.imshow(raw_image)\n",
    "plt.show()\n",
    "raw_image = raw_image.convert(\"RGB\")\n",
    "image = vis_processors[\"eval\"](raw_image).unsqueeze(0)\n",
    "image = image.to(device)\n",
    "print(image.shape)\n",
    "\n",
    "qu = \"Please describe this image in detail.\"\n",
    "# qu = \"What can you see in this image?\"\n",
    "# qu = \"Introduce about this image.\"\n",
    "# qu = \"Is there a bottle in the image?\"\n",
    "template = INSTRUCTION_TEMPLATE[args.model]\n",
    "qu = template.replace(\"<question>\", qu)\n",
    "\n",
    "\n",
    "with torch.inference_mode():\n",
    "    with torch.no_grad():\n",
    "        out = model.generate(\n",
    "            {\"image\": norm(image), \"prompt\":qu}, \n",
    "            use_nucleus_sampling=False, \n",
    "            num_beams=1,\n",
    "            max_new_tokens=512,\n",
    "            output_attentions=True,\n",
    "            # ours_decoding=True,\n",
    "            # scale_factor=50,\n",
    "            # threshold=15.0,\n",
    "            # num_attn_candidates=5,\n",
    "        )\n",
    "print(out[0])\n",
    "print(\"\\n\")\n",
    "\n",
    "\n",
    "qu_append = out[0]\n",
    "qu = qu + qu_append\n",
    "\n",
    "\n",
    "with torch.inference_mode():\n",
    "    with torch.no_grad():\n",
    "        out = model.generate(\n",
    "            {\"image\": norm(image), \"prompt\":qu}, \n",
    "            use_nucleus_sampling=False, \n",
    "            num_beams=1,\n",
    "            max_new_tokens=512,\n",
    "            output_attentions=True,\n",
    "            # ours_decoding=True,\n",
    "            # scale_factor=50,\n",
    "            # threshold=15.0,\n",
    "            # num_attn_candidates=5,\n",
    "            output_attentions=True,\n",
    "            return_dict_in_generate=True,\n",
    "        )\n",
    "\n",
    "# print(len(out[-1]))\n",
    "# print(out.attentions[0].shape)  \n",
    "\n",
    "# print(torch.topk(out.attentions[31].mean(1).squeeze().mean(0), k=10))\n",
    "\n",
    "attns = [attn.clone() for attn in out.attentions]\n",
    "# [layer_id, bs, head_id, qk, v]\n",
    "\n",
    "# for i in range(len(attns)):\n",
    "#     for j in range(len(attns[i])):\n",
    "#         attns[i][j] = attns[i][j] * 5\n",
    "#         attns[i][j] = attns[i][j].clamp(0, 1)\n",
    "\n",
    "\n",
    "try:\n",
    "    p_before, p_after = qu.split('<ImageHere>')\n",
    "    p_before = model.system_message + p_before if args.model in [\"shikra\", \"llava-1.5\"] else p_before\n",
    "    p_before_tokens = model.llama_tokenizer(\n",
    "        p_before, return_tensors=\"pt\", add_special_tokens=False).to(\"cuda\").input_ids\n",
    "    p_after_tokens = model.llama_tokenizer(\n",
    "        p_after, return_tensors=\"pt\", add_special_tokens=False).to(\"cuda\").input_ids\n",
    "    p_before_tokens = model.llama_tokenizer.convert_ids_to_tokens(p_before_tokens[0].tolist())\n",
    "    p_after_tokens = model.llama_tokenizer.convert_ids_to_tokens(p_after_tokens[0].tolist())\n",
    "\n",
    "    bos = torch.ones([1, 1], dtype=torch.int64, device=\"cuda\") * model.llama_tokenizer.bos_token_id\n",
    "    bos_tokens = model.llama_tokenizer.convert_ids_to_tokens(bos[0].tolist())\n",
    "except:\n",
    "    p_before, p_after = qu.split('<ImageHere>')\n",
    "    p_before = model.system_message + p_before if args.model in [\"shikra\", \"llava-1.5\"] else p_before\n",
    "    p_before_tokens = model.llm_tokenizer(\n",
    "        p_before, return_tensors=\"pt\", add_special_tokens=False).to(\"cuda\").input_ids\n",
    "    p_after_tokens = model.llm_tokenizer(\n",
    "        p_after, return_tensors=\"pt\", add_special_tokens=False).to(\"cuda\").input_ids\n",
    "    p_before_tokens = model.llm_tokenizer.convert_ids_to_tokens(p_before_tokens[0].tolist())\n",
    "    p_after_tokens = model.llm_tokenizer.convert_ids_to_tokens(p_after_tokens[0].tolist())\n",
    "\n",
    "    bos = torch.ones([1, 1], dtype=torch.int64, device=\"cuda\") * model.llm_tokenizer.bos_token_id\n",
    "    bos_tokens = model.llm_tokenizer.convert_ids_to_tokens(bos[0].tolist())\n",
    "\n",
    "\n",
    "\n",
    "print(qu)\n",
    "# print(p_before_tokens)\n",
    "# print(p_after_tokens)\n",
    "# print(bos_tokens)\n",
    "\n",
    "NUM_IMAGE_TOKENS = 256 if args.model == \"shikra\" else 32\n",
    "NUM_IMAGE_TOKENS = 576 if args.model == \"llava-1.5\" else NUM_IMAGE_TOKENS\n",
    "tokens = bos_tokens + p_before_tokens + ['img_token'] * NUM_IMAGE_TOKENS + p_after_tokens\n",
    "seq_len = len(tokens)\n",
    "len1 = len(bos_tokens + p_before_tokens)\n",
    "if args.model == \"instructblip\":\n",
    "    len2 = len(bos_tokens + p_before_tokens + ['img_token'] * NUM_IMAGE_TOKENS + p_after_tokens[:p_after_tokens.index(\".\")+1])\n",
    "else:\n",
    "    len2 = len(bos_tokens + p_before_tokens + ['img_token'] * NUM_IMAGE_TOKENS + p_after_tokens[:p_after_tokens.index(\":\")+1])\n",
    "# print(len(tokens))\n",
    "print(len1, len2)\n",
    "\n",
    "tokens = [str(idx) + \"-\" + token for idx, token in enumerate(tokens)]\n",
    "print(tokens)\n",
    "\n",
    "attn_last = attns[-1].max(1).values.data.squeeze()\n",
    "attn_last = attn_last / attn_last.sum(-1, keepdim=True)\n",
    "attn_max = torch.cat([attn.unsqueeze(0) for attn in attns], dim=0).max(2).values.data.max(0).values.data.squeeze()\n",
    "attn_max = attn_max / attn_max.sum(-1, keepdim=True)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "path = \"text_feat/self_attention/{}/\".format(args.model)\n",
    "if not os.path.exists(path):\n",
    "    os.mkdir(path)\n",
    "\n",
    "plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号       \n",
    "plt.rcParams['xtick.direction'] = 'in'\n",
    "# plt.rcParams['ytick.direction'] = 'in'\n",
    "def draw(data, x, y, ax):\n",
    "    seaborn.heatmap(data, \n",
    "                    xticklabels=x, square=True, yticklabels=y, vmin=0, vmax=1.0, \n",
    "                    cbar=False, ax=ax)\n",
    "fig, axs = plt.subplots(1, 1, figsize=(100, 100))#布置画板\n",
    "draw(5*attn_last.cpu().numpy(), x=tokens, y=tokens, ax=axs)\n",
    "plt.show()\n",
    "# plt.savefig(\"xxx.pdf\", dpi=1600)"
   ]
  }
 ],
 "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}