video_chat2/mvbench.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.config import Config\n",
    "config_file = \"configs/config.json\"\n",
    "cfg = Config.from_file(config_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import io\n",
    "import json\n",
    "\n",
    "from models.videochat2_it import VideoChat2_it\n",
    "from utils.easydict import EasyDict\n",
    "import torch\n",
    "\n",
    "from transformers import StoppingCriteria, StoppingCriteriaList\n",
    "\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "import numpy as np\n",
    "from decord import VideoReader, cpu\n",
    "import torchvision.transforms as T\n",
    "from dataset.video_transforms import (\n",
    "    GroupNormalize, GroupScale, GroupCenterCrop, \n",
    "    Stack, ToTorchFormatTensor\n",
    ")\n",
    "from torchvision.transforms.functional import InterpolationMode\n",
    "from torch.utils.data import Dataset\n",
    "\n",
    "from torchvision import transforms\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "from IPython.display import Video, HTML\n",
    "\n",
    "from peft import get_peft_model, LoraConfig, TaskType\n",
    "import copy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.046282052993774414,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": null,
       "postfix": null,
       "prefix": "Loading checkpoint shards",
       "rate": null,
       "total": 2,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6914e9f461564316b49a709622a83ac8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# load stage2 model\n",
    "cfg.model.vision_encoder.num_frames = 4\n",
    "model = VideoChat2_it(config=cfg.model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = model.to(torch.device(cfg.device))\n",
    "model = model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "_IncompatibleKeys(missing_keys=['llama_model.base_model.model.model.embed_tokens.weight', 'llama_model.base_model.model.model.layers.0.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.0.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.0.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.0.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.0.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.0.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.0.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.0.input_layernorm.weight', 'llama_model.base_model.model.model.layers.0.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.1.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.1.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.1.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.1.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.1.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.1.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.1.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.1.input_layernorm.weight', 'llama_model.base_model.model.model.layers.1.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.2.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.2.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.2.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.2.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.2.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.2.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.2.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.2.input_layernorm.weight', 'llama_model.base_model.model.model.layers.2.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.3.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.3.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.3.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.3.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.3.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.3.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.3.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.3.input_layernorm.weight', 'llama_model.base_model.model.model.layers.3.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.4.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.4.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.4.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.4.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.4.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.4.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.4.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.4.input_layernorm.weight', 'llama_model.base_model.model.model.layers.4.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.5.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.5.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.5.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.5.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.5.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.5.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.5.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.5.input_layernorm.weight', 'llama_model.base_model.model.model.layers.5.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.6.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.6.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.6.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.6.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.6.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.6.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.6.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.6.input_layernorm.weight', 'llama_model.base_model.model.model.layers.6.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.7.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.7.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.7.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.7.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.7.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.7.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.7.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.7.input_layernorm.weight', 'llama_model.base_model.model.model.layers.7.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.8.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.8.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.8.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.8.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.8.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.8.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.8.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.8.input_layernorm.weight', 'llama_model.base_model.model.model.layers.8.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.9.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.9.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.9.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.9.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.9.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.9.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.9.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.9.input_layernorm.weight', 'llama_model.base_model.model.model.layers.9.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.10.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.10.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.10.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.10.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.10.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.10.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.10.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.10.input_layernorm.weight', 'llama_model.base_model.model.model.layers.10.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.11.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.11.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.11.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.11.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.11.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.11.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.11.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.11.input_layernorm.weight', 'llama_model.base_model.model.model.layers.11.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.12.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.12.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.12.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.12.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.12.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.12.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.12.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.12.input_layernorm.weight', 'llama_model.base_model.model.model.layers.12.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.13.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.13.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.13.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.13.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.13.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.13.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.13.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.13.input_layernorm.weight', 'llama_model.base_model.model.model.layers.13.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.14.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.14.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.14.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.14.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.14.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.14.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.14.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.14.input_layernorm.weight', 'llama_model.base_model.model.model.layers.14.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.15.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.15.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.15.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.15.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.15.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.15.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.15.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.15.input_layernorm.weight', 'llama_model.base_model.model.model.layers.15.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.16.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.16.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.16.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.16.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.16.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.16.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.16.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.16.input_layernorm.weight', 'llama_model.base_model.model.model.layers.16.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.17.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.17.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.17.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.17.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.17.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.17.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.17.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.17.input_layernorm.weight', 'llama_model.base_model.model.model.layers.17.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.18.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.18.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.18.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.18.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.18.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.18.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.18.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.18.input_layernorm.weight', 'llama_model.base_model.model.model.layers.18.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.19.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.19.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.19.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.19.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.19.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.19.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.19.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.19.input_layernorm.weight', 'llama_model.base_model.model.model.layers.19.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.20.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.20.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.20.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.20.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.20.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.20.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.20.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.20.input_layernorm.weight', 'llama_model.base_model.model.model.layers.20.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.21.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.21.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.21.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.21.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.21.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.21.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.21.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.21.input_layernorm.weight', 'llama_model.base_model.model.model.layers.21.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.22.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.22.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.22.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.22.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.22.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.22.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.22.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.22.input_layernorm.weight', 'llama_model.base_model.model.model.layers.22.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.23.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.23.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.23.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.23.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.23.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.23.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.23.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.23.input_layernorm.weight', 'llama_model.base_model.model.model.layers.23.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.24.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.24.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.24.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.24.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.24.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.24.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.24.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.24.input_layernorm.weight', 'llama_model.base_model.model.model.layers.24.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.25.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.25.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.25.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.25.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.25.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.25.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.25.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.25.input_layernorm.weight', 'llama_model.base_model.model.model.layers.25.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.26.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.26.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.26.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.26.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.26.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.26.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.26.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.26.input_layernorm.weight', 'llama_model.base_model.model.model.layers.26.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.27.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.27.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.27.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.27.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.27.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.27.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.27.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.27.input_layernorm.weight', 'llama_model.base_model.model.model.layers.27.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.28.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.28.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.28.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.28.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.28.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.28.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.28.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.28.input_layernorm.weight', 'llama_model.base_model.model.model.layers.28.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.29.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.29.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.29.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.29.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.29.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.29.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.29.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.29.input_layernorm.weight', 'llama_model.base_model.model.model.layers.29.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.30.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.30.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.30.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.30.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.30.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.30.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.30.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.30.input_layernorm.weight', 'llama_model.base_model.model.model.layers.30.post_attention_layernorm.weight', 'llama_model.base_model.model.model.layers.31.self_attn.q_proj.weight', 'llama_model.base_model.model.model.layers.31.self_attn.k_proj.weight', 'llama_model.base_model.model.model.layers.31.self_attn.v_proj.weight', 'llama_model.base_model.model.model.layers.31.self_attn.o_proj.weight', 'llama_model.base_model.model.model.layers.31.mlp.gate_proj.weight', 'llama_model.base_model.model.model.layers.31.mlp.down_proj.weight', 'llama_model.base_model.model.model.layers.31.mlp.up_proj.weight', 'llama_model.base_model.model.model.layers.31.input_layernorm.weight', 'llama_model.base_model.model.model.layers.31.post_attention_layernorm.weight', 'llama_model.base_model.model.model.norm.weight', 'llama_model.base_model.model.lm_head.weight'], unexpected_keys=[])\n"
     ]
    }
   ],
   "source": [
    "# add lora to run stage3 model\n",
    "peft_config = LoraConfig(\n",
    "    task_type=TaskType.CAUSAL_LM, inference_mode=False, \n",
    "    r=16, lora_alpha=32, lora_dropout=0.\n",
    ")\n",
    "model.llama_model = get_peft_model(model.llama_model, peft_config)\n",
    "\n",
    "state_dict = torch.load(\"your_model_path/videochat2_7b_stage3.pth\", \"cpu\")\n",
    "\n",
    "if 'model' in state_dict.keys():\n",
    "    msg = model.load_state_dict(state_dict['model'], strict=False)\n",
    "else:\n",
    "    msg = model.load_state_dict(state_dict, strict=False)\n",
    "print(msg)\n",
    "\n",
    "model = model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_prompt(conv):\n",
    "    ret = conv.system + conv.sep\n",
    "    for role, message in conv.messages:\n",
    "        if message:\n",
    "            ret += role + \": \" + message + conv.sep\n",
    "        else:\n",
    "            ret += role + \":\"\n",
    "    return ret\n",
    "\n",
    "\n",
    "def get_prompt2(conv):\n",
    "    ret = conv.system + conv.sep\n",
    "    count = 0\n",
    "    for role, message in conv.messages:\n",
    "        count += 1\n",
    "        if count == len(conv.messages):\n",
    "            ret += role + \": \" + message\n",
    "        else:\n",
    "            if message:\n",
    "                ret += role + \": \" + message + conv.sep\n",
    "            else:\n",
    "                ret += role + \":\"\n",
    "    return ret\n",
    "\n",
    "\n",
    "def get_context_emb(conv, model, img_list, answer_prompt=None, print_res=False):\n",
    "    if answer_prompt:\n",
    "        prompt = get_prompt2(conv)\n",
    "    else:\n",
    "        prompt = get_prompt(conv)\n",
    "    if print_res:\n",
    "        print(prompt)\n",
    "    if '<VideoHere>' in prompt:\n",
    "        prompt_segs = prompt.split('<VideoHere>')\n",
    "    else:\n",
    "        prompt_segs = prompt.split('<ImageHere>')\n",
    "    assert len(prompt_segs) == len(img_list) + 1, \"Unmatched numbers of image placeholders and images.\"\n",
    "    with torch.no_grad():\n",
    "        seg_tokens = [\n",
    "            model.llama_tokenizer(\n",
    "                seg, return_tensors=\"pt\", add_special_tokens=i == 0).to(\"cuda:0\").input_ids\n",
    "            # only add bos to the first seg\n",
    "            for i, seg in enumerate(prompt_segs)\n",
    "        ]\n",
    "        seg_embs = [model.llama_model.base_model.model.model.embed_tokens(seg_t) for seg_t in seg_tokens]\n",
    "    mixed_embs = [emb for pair in zip(seg_embs[:-1], img_list) for emb in pair] + [seg_embs[-1]]\n",
    "    mixed_embs = torch.cat(mixed_embs, dim=1)\n",
    "    return mixed_embs\n",
    "\n",
    "\n",
    "def ask(text, conv):\n",
    "    conv.messages.append([conv.roles[0], text + '\\n'])\n",
    "        \n",
    "\n",
    "class StoppingCriteriaSub(StoppingCriteria):\n",
    "    def __init__(self, stops=[], encounters=1):\n",
    "        super().__init__()\n",
    "        self.stops = stops\n",
    "    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor):\n",
    "        for stop in self.stops:\n",
    "            if torch.all((stop == input_ids[0][-len(stop):])).item():\n",
    "                return True\n",
    "        return False\n",
    "    \n",
    "    \n",
    "def answer(conv, model, img_list, do_sample=True, max_new_tokens=200, num_beams=1, min_length=1, top_p=0.9,\n",
    "               repetition_penalty=1.0, length_penalty=1, temperature=1.0, answer_prompt=None, print_res=False):\n",
    "    stop_words_ids = [\n",
    "        torch.tensor([835]).to(\"cuda:0\"),\n",
    "        torch.tensor([2277, 29937]).to(\"cuda:0\")]  # '###' can be encoded in two different ways.\n",
    "    stopping_criteria = StoppingCriteriaList([StoppingCriteriaSub(stops=stop_words_ids)])\n",
    "    \n",
    "    conv.messages.append([conv.roles[1], answer_prompt])\n",
    "    embs = get_context_emb(conv, model, img_list, answer_prompt=answer_prompt, print_res=print_res)\n",
    "    with torch.no_grad():\n",
    "        outputs = model.llama_model.generate(\n",
    "            inputs_embeds=embs,\n",
    "            max_new_tokens=max_new_tokens,\n",
    "            stopping_criteria=stopping_criteria,\n",
    "            num_beams=num_beams,\n",
    "            do_sample=do_sample,\n",
    "            min_length=min_length,\n",
    "            top_p=top_p,\n",
    "            repetition_penalty=repetition_penalty,\n",
    "            length_penalty=length_penalty,\n",
    "            temperature=temperature,\n",
    "        )\n",
    "    output_token = outputs[0]\n",
    "    if output_token[0] == 0:  # the model might output a unknow token <unk> at the beginning. remove it\n",
    "            output_token = output_token[1:]\n",
    "    if output_token[0] == 1:  # some users find that there is a start token <s> at the beginning. remove it\n",
    "            output_token = output_token[1:]\n",
    "    output_text = model.llama_tokenizer.decode(output_token, add_special_tokens=False)\n",
    "    output_text = output_text.split('###')[0]  # remove the stop sign '###'\n",
    "    output_text = output_text.split('Assistant:')[-1].strip()\n",
    "    conv.messages[-1][1] = output_text\n",
    "    return output_text, output_token.cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_index(num_frames, num_segments):\n",
    "    seg_size = float(num_frames - 1) / num_segments\n",
    "    start = int(seg_size / 2)\n",
    "    offsets = np.array([\n",
    "        start + int(np.round(seg_size * idx)) for idx in range(num_segments)\n",
    "    ])\n",
    "    return offsets\n",
    "\n",
    "\n",
    "def load_video(video_path, num_segments=8, return_msg=False, resolution=224):\n",
    "    vr = VideoReader(video_path, ctx=cpu(0), num_threads=1)\n",
    "    num_frames = len(vr)\n",
    "    frame_indices = get_index(num_frames, num_segments)\n",
    "\n",
    "    # transform\n",
    "    crop_size = resolution\n",
    "    scale_size = resolution\n",
    "    input_mean = [0.48145466, 0.4578275, 0.40821073]\n",
    "    input_std = [0.26862954, 0.26130258, 0.27577711]\n",
    "\n",
    "    transform = T.Compose([\n",
    "        GroupScale(int(scale_size), interpolation=InterpolationMode.BICUBIC),\n",
    "        GroupCenterCrop(crop_size),\n",
    "        Stack(),\n",
    "        ToTorchFormatTensor(),\n",
    "        GroupNormalize(input_mean, input_std) \n",
    "    ])\n",
    "\n",
    "    images_group = list()\n",
    "    for frame_index in frame_indices:\n",
    "        img = Image.fromarray(vr[frame_index].numpy())\n",
    "        images_group.append(img)\n",
    "    torch_imgs = transform(images_group)\n",
    "    if return_msg:\n",
    "        fps = float(vr.get_avg_fps())\n",
    "        sec = \", \".join([str(round(f / fps, 1)) for f in frame_indices])\n",
    "        # \" \" should be added in the start and end\n",
    "        msg = f\"The video contains {len(frame_indices)} frames sampled at {sec} seconds.\"\n",
    "        return torch_imgs, msg\n",
    "    else:\n",
    "        return torch_imgs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_sinusoid_encoding_table(n_position=784, d_hid=1024, cur_frame=8, ckpt_num_frame=4, pre_n_position=784): \n",
    "    ''' Sinusoid position encoding table ''' \n",
    "    # TODO: make it with torch instead of numpy \n",
    "    def get_position_angle_vec(position): \n",
    "        return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)] \n",
    "    \n",
    "    # generate checkpoint position embedding\n",
    "    sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(pre_n_position)]) \n",
    "    sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2]) # dim 2i \n",
    "    sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2]) # dim 2i+1 \n",
    "    sinusoid_table = torch.tensor(sinusoid_table, dtype=torch.float, requires_grad=False).unsqueeze(0)\n",
    "    \n",
    "    print(f\"n_position: {n_position}\")\n",
    "    print(f\"pre_n_position: {pre_n_position}\")\n",
    "    \n",
    "    if n_position != pre_n_position:\n",
    "        T = ckpt_num_frame # checkpoint frame\n",
    "        P = 14 # checkpoint size\n",
    "        C = d_hid\n",
    "        new_P = int((n_position // cur_frame) ** 0.5) # testing size\n",
    "        if new_P != 14:\n",
    "            print(f'Pretraining uses 14x14, but current version is {new_P}x{new_P}')\n",
    "            print(f'Interpolate the position embedding')\n",
    "            sinusoid_table = sinusoid_table.reshape(-1, T, P, P, C)\n",
    "            sinusoid_table = sinusoid_table.reshape(-1, P, P, C).permute(0, 3, 1, 2)\n",
    "            sinusoid_table = torch.nn.functional.interpolate(\n",
    "                sinusoid_table, size=(new_P, new_P), mode='bicubic', align_corners=False)\n",
    "            # BT, C, H, W -> BT, H, W, C ->  B, T, H, W, C\n",
    "            sinusoid_table = sinusoid_table.permute(0, 2, 3, 1).reshape(-1, T, new_P, new_P, C)\n",
    "            sinusoid_table = sinusoid_table.flatten(1, 3)  # B, THW, C\n",
    "    \n",
    "    if cur_frame != ckpt_num_frame:\n",
    "        print(f'Pretraining uses 4 frames, but current frame is {cur_frame}')\n",
    "        print(f'Interpolate the position embedding')\n",
    "        T = ckpt_num_frame # checkpoint frame\n",
    "        new_T = cur_frame # testing frame\n",
    "        # interpolate\n",
    "        P = int((n_position // cur_frame) ** 0.5) # testing size\n",
    "        C = d_hid\n",
    "        sinusoid_table = sinusoid_table.reshape(-1, T, P, P, C)\n",
    "        sinusoid_table = sinusoid_table.permute(0, 2, 3, 4, 1).reshape(-1, C, T)  # BHW, C, T\n",
    "        sinusoid_table = torch.nn.functional.interpolate(sinusoid_table, size=new_T, mode='linear')\n",
    "        sinusoid_table = sinusoid_table.reshape(1, P, P, C, new_T).permute(0, 4, 1, 2, 3) # B, T, H, W, C\n",
    "        sinusoid_table = sinusoid_table.flatten(1, 3)  # B, THW, C\n",
    "        \n",
    "    return sinusoid_table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "n_position: 3136\n",
      "pre_n_position: 784\n",
      "Pretraining uses 4 frames, but current frame is 16\n",
      "Interpolate the position embedding\n",
      "The video contains 16 frames sampled at 0.3, 0.9, 1.5, 2.2, 2.8, 3.4, 4.0, 4.7, 5.3, 5.9, 6.5, 7.2, 7.8, 8.4, 9.0, 9.6 seconds.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<video alt=\"test\" controls><source src=\"./example/yoga.mp4\" type=\"video/mp4\"></video>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vid_path = \"./example/yoga.mp4\"\n",
    "# vid_path = \"./example/jesse_dance.mp4\"\n",
    "\n",
    "\n",
    "# num_frame = 8\n",
    "num_frame = 16\n",
    "# resolution = 384\n",
    "resolution = 224\n",
    "vid, msg = load_video(vid_path, num_segments=num_frame, return_msg=True, resolution=resolution)\n",
    "new_pos_emb = get_sinusoid_encoding_table(n_position=(resolution//16)**2*num_frame, cur_frame=num_frame)\n",
    "model.vision_encoder.encoder.pos_embed = new_pos_emb\n",
    "\n",
    "print(msg)\n",
    "    \n",
    "# The model expects inputs of shape: T x C x H x W\n",
    "TC, H, W = vid.shape\n",
    "video = vid.reshape(1, TC//3, 3, H, W).to(\"cuda:0\")\n",
    "\n",
    "img_list = []\n",
    "with torch.no_grad():\n",
    "    image_emb, _ = model.encode_img(video, \"Watch the video and answer the question.\")\n",
    "#     image_emb, _ = model.encode_img(video, \"\")\n",
    "\n",
    "img_list.append(image_emb)\n",
    "\n",
    "HTML(f'<video alt=\"test\" controls><source src=\"{vid_path}\" type=\"video/mp4\"></video>')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "###Human: <Video><VideoHere></Video>\n",
      "###Human: Describe the video in details.\n",
      "###Assistant:\n",
      "The video shows a woman performing yoga on a mat in a beautiful outdoor setting. She is wearing a black tank top and black leggings. The woman is seen performing various yoga poses, including a forward bend, a downward facing dog, and a bridge pose. The video also shows the woman stretching and doing some breathing exercises. The outdoor setting provides a peaceful and serene atmosphere for the yoga practice.\n"
     ]
    }
   ],
   "source": [
    "chat = EasyDict({\n",
    "    \"system\": \"\",\n",
    "    \"roles\": (\"Human\", \"Assistant\"),\n",
    "    \"messages\": [],\n",
    "    \"sep\": \"###\"\n",
    "})\n",
    "\n",
    "chat.messages.append([chat.roles[0], f\"<Video><VideoHere></Video>\\n\"])\n",
    "ask(\"Describe the video in details.\", chat)\n",
    "\n",
    "llm_message = answer(conv=chat, model=model, do_sample=False, img_list=img_list, max_new_tokens=256, print_res=True)[0]\n",
    "print(llm_message)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_list = {\n",
    "    \"Action Sequence\": (\"action_sequence.json\", \"your_data_path/star/Charades_v1_480/\", \"video\", True), # has start & end\n",
    "    \"Action Prediction\": (\"action_prediction.json\", \"your_data_path/star/Charades_v1_480/\", \"video\", True), # has start & end\n",
    "    \"Action Antonym\": (\"action_antonym.json\", \"your_data_path/ssv2_video/\", \"video\", False),\n",
    "    \"Fine-grained Action\": (\"fine_grained_action.json\", \"pyour_data_path/Moments_in_Time_Raw/videos/\", \"video\", False),\n",
    "    \"Unexpected Action\": (\"unexpected_action.json\", \"your_data_path/FunQA_test/test/\", \"video\", False),\n",
    "    \"Object Existence\": (\"object_existence.json\", \"your_data_path/clevrer/video_validation/\", \"video\", False),\n",
    "    \"Object Interaction\": (\"object_interaction.json\", \"your_data_path/star/Charades_v1_480/\", \"video\", True), # has start & end\n",
    "    \"Object Shuffle\": (\"object_shuffle.json\", \"your_data_path/perception/videos/\", \"video\", False),\n",
    "    \"Moving Direction\": (\"moving_direction.json\", \"your_data_path/clevrer/video_validation/\", \"video\", False),\n",
    "    \"Action Localization\": (\"action_localization.json\", \"your_data_path/sta/sta_video/\", \"video\", True),  # has start & end\n",
    "    \"Scene Transition\": (\"scene_transition.json\", \"your_data_path/scene_qa/video/\", \"video\", False),\n",
    "    \"Action Count\": (\"action_count.json\", \"your_data_path/perception/videos/\", \"video\", False),\n",
    "    \"Moving Count\": (\"moving_count.json\", \"your_data_path/clevrer/video_validation/\", \"video\", False),\n",
    "    \"Moving Attribute\": (\"moving_attribute.json\", \"your_data_path/clevrer/video_validation/\", \"video\", False),\n",
    "    \"State Change\": (\"state_change.json\", \"your_data_path/perception/videos/\", \"video\", False),\n",
    "    \"Fine-grained Pose\": (\"fine_grained_pose.json\", \"your_data_path/nturgbd/\", \"video\", False),\n",
    "    \"Character Order\": (\"character_order.json\", \"your_data_path/perception/videos/\", \"video\", False),\n",
    "    \"Egocentric Navigation\": (\"egocentric_navigation.json\", \"your_data_path/vlnqa/\", \"video\", False),\n",
    "    \"Episodic Reasoning\": (\"episodic_reasoning.json\", \"your_data_path/tvqa/frames_fps3_hq/\", \"frame\", True),  # has start & end, read frame\n",
    "    \"Counterfactual Inference\": (\"counterfactual_inference.json\", \"your_data_path/clevrer/video_validation/\", \"video\", False),\n",
    "}\n",
    "\n",
    "data_dir = \"your_mvpbench_path/json\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MVBench_dataset(Dataset):\n",
    "    def __init__(self, data_dir, data_list, num_segments=8, resolution=224):\n",
    "        self.data_list = []\n",
    "        for k, v in data_list.items():\n",
    "            with open(os.path.join(data_dir, v[0]), 'r') as f:\n",
    "                json_data = json.load(f)\n",
    "            for data in json_data:\n",
    "                self.data_list.append({\n",
    "                    'task_type': k,\n",
    "                    'prefix': v[1],\n",
    "                    'data_type': v[2],\n",
    "                    'bound': v[3],\n",
    "                    'data': data\n",
    "                })\n",
    "        \n",
    "        self.decord_method = {\n",
    "            'video': self.read_video,\n",
    "            'gif': self.read_gif,\n",
    "            'frame': self.read_frame,\n",
    "        }\n",
    "        \n",
    "        self.num_segments = num_segments\n",
    "        \n",
    "        # transform\n",
    "        crop_size = resolution\n",
    "        scale_size = resolution\n",
    "        input_mean = [0.48145466, 0.4578275, 0.40821073]\n",
    "        input_std = [0.26862954, 0.26130258, 0.27577711]\n",
    "        self.transform = T.Compose([\n",
    "            GroupScale(int(scale_size), interpolation=InterpolationMode.BICUBIC),\n",
    "            GroupCenterCrop(crop_size),\n",
    "            Stack(),\n",
    "            ToTorchFormatTensor(),\n",
    "            GroupNormalize(input_mean, input_std) \n",
    "        ])\n",
    "    \n",
    "    def __str__(self):\n",
    "        len_list = {}\n",
    "        option_list = {}\n",
    "        for data in self.data_list:\n",
    "            if data['task_type'] not in len_list:\n",
    "                len_list[data['task_type']] = 0\n",
    "            len_list[data['task_type']] += 1\n",
    "            if data['task_type'] not in option_list:\n",
    "                option_list[data['task_type']] = 0\n",
    "            option_list[data['task_type']] += len(data['data']['candidates'])\n",
    "        \n",
    "        correct = 0\n",
    "        total = 0\n",
    "        res = f\"There are {len(self.data_list)} videos as follow:\\n\"\n",
    "        for k, v in len_list.items():\n",
    "            correct += len_list[k]\n",
    "            total += option_list[k]\n",
    "            res += f\"{v} for {k} ({option_list[k]} options => {len_list[k]/option_list[k]*100:.2f}%)\\n\"\n",
    "            correct = correct + 1 / option_list[k]\n",
    "        res += f\"Total random accuracy: {correct/total*100:.2f}%\"\n",
    "        return res.rstrip()\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.data_list)\n",
    "    \n",
    "    def get_index(self, bound, fps, max_frame, first_idx=0):\n",
    "        if bound:\n",
    "            start, end = bound[0], bound[1]\n",
    "        else:\n",
    "            start, end = -100000, 100000\n",
    "        start_idx = max(first_idx, round(start * fps))\n",
    "        end_idx = min(round(end * fps), max_frame)\n",
    "        seg_size = float(end_idx - start_idx) / self.num_segments\n",
    "        frame_indices = np.array([\n",
    "            int(start_idx + (seg_size / 2) + np.round(seg_size * idx))\n",
    "            for idx in range(self.num_segments)\n",
    "        ])\n",
    "        return frame_indices\n",
    "    \n",
    "    def read_video(self, video_path, bound=None):\n",
    "        vr = VideoReader(video_path, ctx=cpu(0), num_threads=1)\n",
    "        max_frame = len(vr) - 1\n",
    "        fps = float(vr.get_avg_fps())\n",
    "        \n",
    "        images_group = list()\n",
    "        frame_indices = self.get_index(bound, fps, max_frame, first_idx=0) \n",
    "        for frame_index in frame_indices:\n",
    "            img = Image.fromarray(vr[frame_index].numpy())\n",
    "            images_group.append(img)\n",
    "        torch_imgs = self.transform(images_group)\n",
    "        return torch_imgs\n",
    "    \n",
    "    def read_gif(self, video_path, bound=None, fps=25):\n",
    "        gif = imageio.get_reader(video_path)\n",
    "        max_frame = len(gif) - 1\n",
    "        \n",
    "        images_group = list()\n",
    "        frame_indices = self.get_index(bound, fps, max_frame, first_idx=0) \n",
    "        for index, frame in enumerate(gif):\n",
    "            if index in frame_indices:\n",
    "                img = cv2.cvtColor(frame, cv2.COLOR_RGBA2RGB)\n",
    "                img = Image.fromarray(img)\n",
    "                images_group.append(img)\n",
    "        torch_imgs = self.transform(images_group)\n",
    "        return torch_imgs\n",
    "    \n",
    "    def read_frame(self, video_path, bound=None, fps=3):\n",
    "        max_frame = len(os.listdir(video_path))\n",
    "        images_group = list()\n",
    "        frame_indices = self.get_index(bound, fps, max_frame, first_idx=1) # frame_idx starts from 1\n",
    "        for frame_index in frame_indices:\n",
    "            img = Image.open(os.path.join(video_path, f\"{frame_index:05d}.jpg\"))\n",
    "            images_group.append(img)\n",
    "        torch_imgs = self.transform(images_group)\n",
    "        return torch_imgs\n",
    "\n",
    "    def qa_template(self, data):\n",
    "        question = f\"Question: {data['question']}\\n\"\n",
    "        question += \"Options:\\n\"\n",
    "        answer = data['answer']\n",
    "        answer_idx = -1\n",
    "        for idx, c in enumerate(data['candidates']):\n",
    "            question += f\"({chr(ord('A') + idx)}) {c}\\n\"\n",
    "            if c == answer:\n",
    "                answer_idx = idx\n",
    "        question = question.rstrip()\n",
    "        answer = f\"({chr(ord('A') + answer_idx)}) {answer}\"\n",
    "        return question, answer\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        decord_method = self.decord_method[self.data_list[idx]['data_type']]\n",
    "        bound = None\n",
    "        if self.data_list[idx]['bound']:\n",
    "            bound = (\n",
    "                self.data_list[idx]['data']['start'],\n",
    "                self.data_list[idx]['data']['end'],\n",
    "            )\n",
    "        video_path = os.path.join(self.data_list[idx]['prefix'], self.data_list[idx]['data']['video'])\n",
    "        torch_imgs = decord_method(video_path, bound)\n",
    "        question, answer = self.qa_template(self.data_list[idx]['data'])\n",
    "            \n",
    "        return {\n",
    "            'video': torch_imgs, \n",
    "            'question': question, \n",
    "            'answer': answer,\n",
    "            'task_type': self.data_list[idx]['task_type']\n",
    "        }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "n_position: 3136\n",
      "pre_n_position: 784\n",
      "Pretraining uses 4 frames, but current frame is 16\n",
      "Interpolate the position embedding\n"
     ]
    }
   ],
   "source": [
    "#  position embedding\n",
    "num_frame = 16\n",
    "resolution = 224\n",
    "new_pos_emb = get_sinusoid_encoding_table(n_position=(resolution//16)**2*num_frame, cur_frame=num_frame)\n",
    "model.vision_encoder.encoder.pos_embed = new_pos_emb\n",
    "\n",
    "dataset = MVBench_dataset(data_dir, data_list, num_segments=num_frame, resolution=resolution)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "def infer_mvbench(\n",
    "        data_sample, system=\"\", \n",
    "        question_prompt='', # add in the end of question\n",
    "        answer_prompt=None, # add in the begining of answer\n",
    "        return_prompt='',  # add in the begining of return message\n",
    "        system_q=False, # whether add question in the system prompt for QFormer\n",
    "        print_res=True,\n",
    "        system_llm=False\n",
    "    ):\n",
    "    video = data_sample[\"video\"]\n",
    "    TC, H, W = video.shape\n",
    "    video = video.reshape(1, TC//3, 3, H, W).to(\"cuda:0\")\n",
    "    \n",
    "    video_list = []\n",
    "    with torch.no_grad():\n",
    "        if system_q:\n",
    "            video_emb, _ = model.encode_img(video, system + data_sample['question'])\n",
    "        else:\n",
    "            video_emb, _ = model.encode_img(video, system)\n",
    "    video_list.append(video_emb)\n",
    "#     video_list.append(torch.zeros_like(video_emb))\n",
    "\n",
    "    chat = EasyDict({\n",
    "        \"system\": system,\n",
    "        \"roles\": (\"Human\", \"Assistant\"),\n",
    "        \"messages\": [],\n",
    "        \"sep\": \"###\"\n",
    "    })\n",
    "\n",
    "    chat.messages.append([chat.roles[0], f\"<Video><VideoHere></Video>\\n\"])\n",
    "    \n",
    "    if system_llm:\n",
    "        prompt = system + data_sample['question'] + question_prompt\n",
    "    else:\n",
    "        prompt = data_sample['question'] + question_prompt\n",
    "    \n",
    "    ask(prompt, chat)\n",
    "\n",
    "    llm_message = answer(\n",
    "        conv=chat, model=model, do_sample=False, \n",
    "        img_list=video_list, max_new_tokens=100, \n",
    "        answer_prompt=answer_prompt, print_res=print_res\n",
    "    )[0]\n",
    "    # remove potential explanation\n",
    "    llm_message = return_prompt + llm_message.strip().split('\\n')[0]\n",
    "    print(llm_message)\n",
    "    print(f\"GT: {data_sample['answer']}\")\n",
    "    return llm_message"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_ans(pred, gt):\n",
    "    flag = False\n",
    "    \n",
    "    pred_list = pred.lower().split(' ')\n",
    "    pred_option, pred_content = pred_list[0], ' '.join(pred_list[1:])\n",
    "    gt_list = gt.lower().split(' ')\n",
    "    gt_option, gt_content = gt_list[0], ' '.join(gt_list[1:])\n",
    "    if gt_content[-1] == '.':\n",
    "        gt_content = gt_content[:-1]\n",
    "    \n",
    "    if pred_option.replace('.', '') in gt_option:\n",
    "        flag = True\n",
    "    elif gt_option in pred_option:\n",
    "        flag = True\n",
    "        \n",
    "    return flag"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Carefully watch the video and pay attention to the cause and sequence of events, the detail and movement of objects, and the action and pose of persons. Based on your observations, select the best option that accurately addresses the question.\n",
      "###Human: <Video><VideoHere></Video>\n",
      "###Human: Carefully watch the video and pay attention to the cause and sequence of events, the detail and movement of objects, and the action and pose of persons. Based on your observations, select the best option that accurately addresses the question.\n",
      "Question: What happened after the person took the food?\n",
      "Options:\n",
      "(A) Ate the medicine.\n",
      "(B) Tidied up the blanket.\n",
      "(C) Put down the cup/glass/bottle.\n",
      "(D) Took the box.\n",
      "Only give the best option.\n",
      "###Assistant: Best option:(\n"
     ]
    }
   ],
   "source": [
    "save_path = \"./test\"\n",
    "\n",
    "correct = 0\n",
    "total = 0\n",
    "res_list = []\n",
    "acc_dict = {}\n",
    "\n",
    "for example in tqdm(dataset):\n",
    "    task_type = example['task_type']\n",
    "    if task_type not in acc_dict:\n",
    "        acc_dict[task_type] = [0, 0] # correct, total\n",
    "    acc_dict[task_type][1] += 1\n",
    "    total += 1\n",
    "    pred = infer_mvbench(\n",
    "        example, \n",
    "        system=\"Carefully watch the video and pay attention to the cause and sequence of events, the detail and movement of objects, and the action and pose of persons. Based on your observations, select the best option that accurately addresses the question.\\n\",\n",
    "        question_prompt=\"\\nOnly give the best option.\",\n",
    "        answer_prompt=\"Best option:(\",\n",
    "        return_prompt='(',\n",
    "        system_q=False,\n",
    "        print_res=True,\n",
    "        system_llm=True\n",
    "    )\n",
    "    gt = example['answer']\n",
    "    res_list.append({\n",
    "        'pred': pred,\n",
    "        'gt': gt\n",
    "    })\n",
    "    if check_ans(pred=pred, gt=gt):\n",
    "        acc_dict[task_type][0] += 1\n",
    "        correct += 1\n",
    "    print(f\"Part  Acc: {acc_dict[task_type][0] / acc_dict[task_type][1] * 100 :.2f}%\")\n",
    "    print(f\"Total Acc: {correct / total * 100 :.2f}%\")\n",
    "    print('-' * 30, task_type, '-' * 30)\n",
    "\n",
    "with open(f\"{save_path}.json\", \"w\") as f:\n",
    "    json.dump({\n",
    "        \"acc_dict\": acc_dict,\n",
    "        \"res_list\": res_list\n",
    "    }, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_res = dict()\n",
    "correct = 0\n",
    "total = 0\n",
    "for k, v in acc_dict.items():\n",
    "    final_res[k] = v[0] / v[1] * 100\n",
    "    correct += v[0]\n",
    "    total += v[1]    \n",
    "final_res['Avg'] = correct / total * 100\n",
    "\n",
    "print(final_res)\n",
    "\n",
    "with open(\"upload_leaderboard.json\", \"w\") as f:\n",
    "    json.dump(final_res, f)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vl",
   "language": "python",
   "name": "vl"
  },
  "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.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}