video_chat2/demo/demo_mistral.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.chdir('../')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.config import Config\n",
    "config_file = \"configs/config_mistral.json\"\n",
    "cfg = Config.from_file(config_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import io\n",
    "\n",
    "from models import VideoChat2_it_mistral\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 torchvision.transforms import PILToTensor\n",
    "from torchvision import transforms\n",
    "from dataset.video_transforms import (\n",
    "    GroupNormalize, GroupScale, GroupCenterCrop, \n",
    "    Stack, ToTorchFormatTensor\n",
    ")\n",
    "from torch.utils.data import Dataset\n",
    "from torchvision.transforms.functional import InterpolationMode\n",
    "\n",
    "from torchvision import transforms\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from IPython.display import Video, HTML\n",
    "\n",
    "from peft import get_peft_model, LoraConfig, TaskType\n",
    "import copy\n",
    "\n",
    "import json\n",
    "from collections import OrderedDict\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "import time\n",
    "import decord\n",
    "decord.bridge.set_bridge(\"torch\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load stage2 model\n",
    "cfg.model.vision_encoder.num_frames = 4\n",
    "model = VideoChat2_it_mistral(config=cfg.model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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",
    "    target_modules=[\n",
    "        \"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n",
    "         \"gate_proj\", \"up_proj\", \"down_proj\", \"lm_head\"\n",
    "    ]\n",
    ")\n",
    "model.mistral_model = get_peft_model(model.mistral_model, peft_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "state_dict = torch.load(\"your_model_path/videochat2/videochat2_mistral_7b_stage3.pth\", \"cpu\")\n",
    "\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.to(torch.device(cfg.device))\n",
    "model = model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.mistral_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.mistral_model.base_model.model.model.embed_tokens(seg_t) for seg_t in seg_tokens]\n",
    "#         seg_embs = [model.mistral_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",
    "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([2]).to(\"cuda:0\"),\n",
    "        torch.tensor([29871, 2]).to(\"cuda:0\")]  # '</s>' 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.mistral_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.mistral_tokenizer.decode(output_token, add_special_tokens=False)\n",
    "    output_text = output_text.split('</s>')[0]  # remove the stop sign </s>\n",
    "#     output_text = output_text.split('[/INST]')[-1].strip()\n",
    "    conv.messages[-1][1] = output_text + '</s>'\n",
    "    return output_text, output_token.cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": null,
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vid_path = \"./demo/example/yoga.mp4\"\n",
    "# vid_path = \"./demo/example/jesse_dance.mp4\"\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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chat = EasyDict({\n",
    "    \"system\": \"\",\n",
    "    \"roles\": (\"[INST]\", \"[/INST]\"),\n",
    "    \"messages\": [],\n",
    "    \"sep\": \"\"\n",
    "})\n",
    "\n",
    "chat.messages.append([chat.roles[0], \"<Video><VideoHere></Video> [/INST]\"])\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=512, print_res=True)[0]\n",
    "print(llm_message)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# img_path = \"./demo/example/dog.png\"\n",
    "img_path = \"./demo/example/bear.jpg\"\n",
    "img = Image.open(img_path).convert('RGB')\n",
    "\n",
    "plt.imshow(img)\n",
    "\n",
    "resolution = 224\n",
    "# resolution = 384\n",
    "new_pos_emb = get_sinusoid_encoding_table(n_position=(resolution//16)**2, cur_frame=1, ckpt_num_frame=1, pre_n_position=14*14)\n",
    "model.vision_encoder.encoder.img_pos_embed = new_pos_emb\n",
    "\n",
    "transform = transforms.Compose(\n",
    "    [\n",
    "        transforms.Resize(\n",
    "            (resolution, resolution), interpolation=InterpolationMode.BICUBIC\n",
    "        ),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),\n",
    "    ]\n",
    ")\n",
    "\n",
    "img = transform(img).unsqueeze(0).unsqueeze(0).cuda()\n",
    "img_list = []\n",
    "with torch.no_grad():\n",
    "#     image_emb, _ = model.encode_img(img, \"\")\n",
    "    image_emb, _ = model.encode_img(img, \"Observe the image and answer the question.\")\n",
    "img_list.append(image_emb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chat = EasyDict({\n",
    "    \"system\": \"\",\n",
    "    \"roles\": (\"[INST]\", \"[/INST]\"),\n",
    "    \"messages\": [],\n",
    "    \"sep\": \"\"\n",
    "})\n",
    "\n",
    "chat.messages.append([chat.roles[0], f\"<Image><ImageHere></Image> [/INST]\"])\n",
    "ask(\"Describe the following image 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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_answer_egoschema(pred, qid):\n",
    "    correct = 0\n",
    "    answer_content = ans_dict[qid]['content'].lower()\n",
    "    if answer_content[-1] == \".\":\n",
    "        answer_content = answer_content[:-1]\n",
    "    if ans_dict[qid]['answer'].lower() in pred.lower():\n",
    "        flag = True\n",
    "        for kk in [\"(A)\", \"(B)\", \"(C)\", \"(D)\", \"(E)\"]:\n",
    "            if kk != ans_dict[qid]['answer'].lower() and kk in pred.lower():\n",
    "                flag = ans_dict\n",
    "                break\n",
    "        if flag:\n",
    "            correct += 1\n",
    "    elif answer_content in pred.lower():\n",
    "        correct = 1\n",
    "    elif answer_content.replace(\"a \", \"\") in pred.lower():\n",
    "        correct = 1\n",
    "    elif answer_content.replace(\"an \", \"\") in pred.lower():\n",
    "        correct = 1\n",
    "    return correct\n",
    "\n",
    "def infer_egoschema(\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",
    "        num_segments=8,\n",
    "    ):\n",
    "    vid_path = os.path.join(\"your_data_path/egoschema/videos\", data_sample['video'])\n",
    "    video, _ = load_video(vid_path, num_segments=num_segments, return_msg=True)\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",
    "\n",
    "    chat = EasyDict({\n",
    "        \"system\": system,\n",
    "        \"roles\": (\"[INST]\", \"[/INST]\"),\n",
    "        \"messages\": [],\n",
    "        \"sep\": \"\"\n",
    "    })\n",
    "\n",
    "    chat.messages.append([chat.roles[0], f\"<Video><VideoHere></Video> [/INST]\"])\n",
    "    \n",
    "    if system_llm:\n",
    "        prompt = system + data_sample['QA'][0]['q'] + question_prompt\n",
    "    else:\n",
    "        prompt = data_sample['QA'][0]['q'] + 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['QA'][0]['a']}\")\n",
    "    return llm_message\n",
    "\n",
    "\n",
    "import csv\n",
    "# You can find the csv files in https://github.com/imagegridworth/IG-VLM/blob/main/data/multiple_choice_qa/EgoSchema.csv\n",
    "with open(\"your_data_path/EgoSchema.csv\", mode='r', encoding='utf-8') as file:\n",
    "    reader = csv.reader(file)\n",
    "\n",
    "    json_data = []\n",
    "    ans_dict = {}\n",
    "    \n",
    "    for idx, msg in enumerate(reader):\n",
    "        if idx == 0:\n",
    "            print(msg)\n",
    "            continue\n",
    "            \n",
    "        video = msg[1] + '.mp4'\n",
    "        input_str = f\"Question: {msg[3].capitalize()}\\nOptions:\\n\"\n",
    "    \n",
    "        target_index = -1\n",
    "        for i, candidate in enumerate(msg[5:]):\n",
    "            option = chr(ord('A') + i)\n",
    "            input_str += f\"({option}) {candidate}\\n\"\n",
    "            if candidate == msg[4]:\n",
    "                target_index = i\n",
    "            \n",
    "        assert target_index != -1\n",
    "        correct = chr(ord('A') + target_index)\n",
    "        \n",
    "        json_data.append({\n",
    "            'video': video,\n",
    "            \"QA\": [{\n",
    "                \"i\": \"\",\n",
    "                \"q\": input_str.strip(),\n",
    "                \"a\": f\"Answer: ({correct}) {msg[4]}\",\n",
    "            }]\n",
    "        })\n",
    "\n",
    "        ans_dict[idx - 1] = {\n",
    "            'video': video,\n",
    "            'answer': f\"({correct})\",\n",
    "            'content': msg[4],\n",
    "        }\n",
    "\n",
    "\n",
    "#  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",
    "correct = 0\n",
    "total = 0\n",
    "total_num = len(json_data)\n",
    "\n",
    "output = \"\"\n",
    "\n",
    "for idx, example in enumerate(tqdm(json_data)):\n",
    "    start = time.time()\n",
    "    llm_message = infer_egoschema(\n",
    "        example, \n",
    "        \"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=False,\n",
    "        system_llm=False,\n",
    "        num_segments=16\n",
    "    )\n",
    "    \n",
    "    duration = time.time() - start\n",
    "    output += (example[\"video\"] + '\\n')\n",
    "    output += (llm_message + '\\n')\n",
    "    correct += check_answer_egoschema(llm_message, idx)\n",
    "    total += 1\n",
    "    print(\"Acc:\", correct / total)\n",
    "    print('-' * 20, f'{idx+1}/{total_num} done,', f'cost: {duration:.2f}s', '-' * 20)\n",
    "\n",
    "with open(\"./demo/egoschema/your_prediction.txt\", \"w\") as f:\n",
    "    f.writelines(output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can find the csv files in https://github.com/egoschema/EgoSchema/blob/main/questions.json\n",
    "with open(\"your_data_path/EgoSchema/questions.json\", \"r\") as f:\n",
    "    full_data = json.load(f)\n",
    "\n",
    "full_egoschema = []\n",
    "for data in full_data:\n",
    "    video = data['q_uid'] + '.mp4'\n",
    "    input_str = f\"Question: {data['question'].capitalize()}\\nOptions:\\n\"\n",
    "\n",
    "    for i, candidate in enumerate(['option 0', 'option 1', 'option 2', 'option 3', 'option 4']):\n",
    "        option = chr(ord('A') + i)\n",
    "        input_str += f\"({option}) {data[candidate]}\\n\"\n",
    "    \n",
    "    full_egoschema.append({\n",
    "        'q_uid': data['q_uid'],\n",
    "        'video': video,\n",
    "        \"QA\": [{\n",
    "            \"i\": \"\",\n",
    "            \"q\": input_str.strip(),\n",
    "            \"a\": \"\",\n",
    "        }]\n",
    "    })\n",
    "\n",
    "\n",
    "def infer_full_egoschema(\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",
    "        num_segments=8,\n",
    "    ):\n",
    "    vid_path = os.path.join(\"your_data_path/egoschema/videos\", data_sample['video'])\n",
    "    print(vid_path)\n",
    "    video, _ = load_video(vid_path, num_segments=num_segments, return_msg=True)\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",
    "\n",
    "    chat = EasyDict({\n",
    "        \"system\": system,\n",
    "        \"roles\": (\"[INST]\", \"[/INST]\"),\n",
    "        \"messages\": [],\n",
    "        \"sep\": \"\"\n",
    "    })\n",
    "\n",
    "    chat.messages.append([chat.roles[0], f\"<Video><VideoHere></Video> [/INST]\"])\n",
    "    \n",
    "    if system_llm:\n",
    "        prompt = system + data_sample['QA'][0]['q'] + question_prompt\n",
    "    else:\n",
    "        prompt = data_sample['QA'][0]['q'] + 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",
    "    return llm_message\n",
    "\n",
    "\n",
    "#  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",
    "\n",
    "ans_dict = {}\n",
    "\n",
    "for idx, example in enumerate(tqdm(full_egoschema)):\n",
    "    start = time.time()\n",
    "    llm_message = infer_full_egoschema(\n",
    "        example, \n",
    "        \"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=False,\n",
    "        system_llm=False,\n",
    "        num_segments=16,\n",
    "    )\n",
    "\n",
    "    assert llm_message[0] == '(' and llm_message[2] == ')'\n",
    "    ans = ord(llm_message[1]) - ord('A')\n",
    "    assert ans in [0, 1, 2, 3, 4]\n",
    "    ans_dict[example['q_uid']] = ans\n",
    "\n",
    "\n",
    "with open(\"./demo/egoschema/your_prediction.json\", \"w\") as f:\n",
    "    json.dump(ans_dict, f)\n",
    "\n",
    "# Then you can run https://github.com/egoschema/EgoSchema/blob/main/validate.py to get the score\n",
    "# python3 validate.py --f ./your_prediction.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import webvtt\n",
    "import re\n",
    "\n",
    "def clean_text(text):\n",
    "    cleaned_text = re.sub(r'[^A-Za-z0-9\\s]', '', text)\n",
    "    return cleaned_text\n",
    "\n",
    "\n",
    "def read_vtt_and_concatenate(file_path, tokenizer, max_len=4096):\n",
    "    prev = \"\"\n",
    "    subtitles = []\n",
    "    for caption in webvtt.read(file_path):\n",
    "        # Split the caption text into individual lines\n",
    "        lines = caption.text.split('\\n')\n",
    "        for line in lines:\n",
    "            # Clean the text and check for repetition\n",
    "            line = clean_text(line)\n",
    "            if prev != line and line:\n",
    "                subtitles.append(line)\n",
    "                prev = line\n",
    "\n",
    "    # Join subtitles to check length\n",
    "    full_text = ' '.join(subtitles)\n",
    "    tokenized_ids = tokenizer(full_text, add_special_tokens=False).input_ids\n",
    "\n",
    "    # If the tokenized length is within the limit, return the full text\n",
    "    if len(tokenized_ids) <= max_len:\n",
    "        return full_text\n",
    "\n",
    "    # Otherwise, we need to trim the text to fit within the limit\n",
    "    # We will keep the first half and the last half\n",
    "    half_len = max_len // 2\n",
    "    start_text = ' '.join(subtitles[:half_len])\n",
    "    end_text = ' '.join(subtitles[-half_len:])\n",
    "    \n",
    "    # Re-tokenize to ensure the total length is within the limit\n",
    "    start_tokenized_ids = tokenizer(start_text, add_special_tokens=False).input_ids\n",
    "    end_tokenized_ids = tokenizer(end_text, add_special_tokens=False).input_ids\n",
    "\n",
    "    # Adjust the lengths to fit within the max_len\n",
    "    while len(start_tokenized_ids) + len(end_tokenized_ids) > max_len:\n",
    "        if len(start_tokenized_ids) > len(end_tokenized_ids):\n",
    "            start_tokenized_ids.pop()\n",
    "        else:\n",
    "            end_tokenized_ids.pop(0)\n",
    "    \n",
    "    # Combine the adjusted parts\n",
    "    adjusted_text = tokenizer.decode(start_tokenized_ids) + ' ... ' + tokenizer.decode(end_tokenized_ids)\n",
    "    \n",
    "    return adjusted_text\n",
    "\n",
    "    \n",
    "class MME_dataset(Dataset):\n",
    "    def __init__(self, data_prefix, anno_path, num_segments=16, resolution=224, max_subtitle_len=4096):\n",
    "        self.data_prefix = data_prefix\n",
    "        with open(anno_path, 'r') as f:\n",
    "            self.data_list = json.load(f)\n",
    "            \n",
    "        self.num_segments = num_segments\n",
    "        self.max_subtitle_len = max_subtitle_len\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",
    "        task_dict = {}\n",
    "        total = 0\n",
    "        for data in self.data_list:\n",
    "            if data['duration_category'] not in ans_dict:\n",
    "                task_dict[data['duration_category']] = {}\n",
    "            for q in data['questions']:\n",
    "                if q['task_type'] not in ans_dict[data['duration_category']]:\n",
    "                    ans_dict[data['duration_category']][q['task_type']] = 0\n",
    "                ans_dict[data['duration_category']][q['task_type']] += 1\n",
    "                total += 1\n",
    "\n",
    "        res = f\"There are {len(self.data_list)} videos.\\n\"\n",
    "        res += f\"There are {total} QAs.\\n\"\n",
    "        for k, v in task_dict.items():\n",
    "            res += f\"------{k}------\\n\"\n",
    "            for kk, vv in task_dict.items():\n",
    "                res += f\"{kk}: {vv}\\n\"\n",
    "                \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_frame(self, video_path, bound=None):\n",
    "        video_path = os.path.join(video_path, str(self.num_segments))\n",
    "        \n",
    "        if os.path.exists(video_path):\n",
    "            frame_list = [p for p in os.listdir(video_path)]\n",
    "        else:\n",
    "            raise Exception\n",
    "            \n",
    "        images_group = list()\n",
    "        \n",
    "        for frame_name in frame_list:\n",
    "            img = Image.open(os.path.join(video_path, frame_name))\n",
    "            images_group.append(img)\n",
    "        torch_imgs = self.transform(images_group)\n",
    "        return torch_imgs\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 qa_template(self, data):\n",
    "        question = f\"Question: {data['question']}\\n\"\n",
    "        question += \"Options:\\n\"\n",
    "        answer = data['answer']\n",
    "        answer = f\"({answer}) {data['choices'][ord(answer) - ord('A')][3:]}\"\n",
    "        for idx, c in enumerate(data['choices']):\n",
    "            cur_choice, cur_text = c[0], c[3:]\n",
    "            question += f\"({cur_choice}) {cur_text}\\n\"\n",
    "        question = question.rstrip()\n",
    "        return question, answer\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        video_name = self.data_list[idx]['url'].split(\"watch?v=\")[1]\n",
    "        video_path = os.path.join(self.data_prefix, \"frames\", video_name)\n",
    "\n",
    "        # We store the videos with only 16 or 32 frames for testing,\n",
    "        # since directly reading the whold videos cost a lot of time.\n",
    "        # You can also read the whole video via self.read_video(video_path)\n",
    "        torch_imgs = self.read_frame(video_path)\n",
    "        duration_category = self.data_list[idx]['duration_category']\n",
    "        qa_list = []\n",
    "        for qa in self.data_list[idx]['questions']:\n",
    "            qa_list.append(self.qa_template(qa))\n",
    "\n",
    "        subtitle = \"\"\n",
    "        try:\n",
    "            subtitle_path = os.path.join(self.data_prefix, \"subtitle\", video_name + \".vtt\")\n",
    "            if os.path.exists(subtitle_path):\n",
    "                subtitle = read_vtt_and_concatenate(subtitle_path, model.mistral_tokenizer, self.max_subtitle_len)\n",
    "        except Exception:\n",
    "            subtitle = \"\"\n",
    "            print(f\"Error for {subtitle_path}\")\n",
    "            \n",
    "        return {\n",
    "            'subtitle': subtitle,\n",
    "            'video': torch_imgs, \n",
    "            'qa_list': qa_list,\n",
    "            'duration_category': duration_category\n",
    "        }\n",
    "    \n",
    "\n",
    "def infer_mme(\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",
    "        add_subtitle=False,\n",
    "    ):\n",
    "    assert system_q == False, \"do not support system_q now\"\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",
    "            raise NotImplementedError\n",
    "        else:\n",
    "            video_emb, _ = model.encode_img(video, system)\n",
    "    video_list.append(video_emb)\n",
    "\n",
    "    pred_list = []\n",
    "    gt_list = []\n",
    "    for idx, qa in enumerate(data_sample['qa_list']):\n",
    "        print(f\"----------qa_{idx}---------\", flush=True)\n",
    "        chat = EasyDict({\n",
    "            \"system\": system,\n",
    "            \"roles\": (\"[INST]\", \"[/INST]\"),\n",
    "            \"messages\": [],\n",
    "            \"sep\": \"\"\n",
    "        })\n",
    "    \n",
    "        if add_subtitle:\n",
    "            if data_sample['subtitle'] != '':\n",
    "                subtitle = f\"This video's subtitles are listed below: {data_sample['subtitle']}\"\n",
    "                chat.messages.append([chat.roles[0], f\"{subtitle}\\n<Video><VideoHere></Video> [/INST]\"])\n",
    "            else:\n",
    "                chat.messages.append([chat.roles[0], f\"<Video><VideoHere></Video> [/INST]\"])\n",
    "        else:\n",
    "            chat.messages.append([chat.roles[0], f\"<Video><VideoHere></Video> [/INST]\"])\n",
    "    \n",
    "        if system_llm:\n",
    "            prompt = system + qa[0] + question_prompt\n",
    "        else:\n",
    "            prompt = qa[0] + 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(f\"Pred: {llm_message}\", flush=True)\n",
    "        print(f\"GT: {qa[1]}\", flush=True)\n",
    "        pred_list.append(llm_message[1])\n",
    "        gt_list.append(qa[1][1])\n",
    "    return pred_list, gt_list\n",
    "\n",
    "    \n",
    "#  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",
    "data_dir = \"your_data_path/videomme\"\n",
    "anno_path =  \"your_data_path/Video-MME.json\"\n",
    "dataset = MME_dataset(\n",
    "    data_dir, \n",
    "    anno_path, \n",
    "    num_segments=num_frame, resolution=resolution\n",
    ")\n",
    "\n",
    "with open(anno_path, 'r') as f:\n",
    "    res_json_data = json.load(f)\n",
    "\n",
    "save_path = \"./demo/videomme/your_prediction\"\n",
    "\n",
    "correct = 0\n",
    "total = 0\n",
    "res_list = []\n",
    "acc_dict = {}\n",
    "\n",
    "for idx, example in enumerate(tqdm(dataset)):\n",
    "    duration_category = example['duration_category']\n",
    "    if duration_category not in acc_dict:\n",
    "        acc_dict[duration_category] = [0, 0] # correct, total\n",
    "    qa_count = len(example['qa_list'])\n",
    "    acc_dict[duration_category][1] += qa_count\n",
    "    total += qa_count\n",
    "    pred_list, gt_list = infer_mme(\n",
    "        example, \n",
    "        \"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=False,\n",
    "        system_llm=True,\n",
    "        # add_subtitle=True, # Comment this line to add subtitles, we use the whole subtitles by default.\n",
    "    )\n",
    "    res_list.append({\n",
    "        'pred': pred_list,\n",
    "        'gt': gt_list\n",
    "    })\n",
    "    qa_idx = 0\n",
    "    for pred, gt in zip(pred_list, gt_list):\n",
    "        if pred == gt:\n",
    "            acc_dict[duration_category][0] += 1\n",
    "            correct += 1\n",
    "        res_json_data[idx]['questions'][qa_idx]['response'] = pred\n",
    "        qa_idx += 1\n",
    "    print(f\"Part  Acc: {acc_dict[duration_category][0] / acc_dict[duration_category][1] * 100 :.2f}%\")\n",
    "    print(f\"Total Acc: {correct / total * 100 :.2f}%\")\n",
    "    print('-' * 50, duration_category, '-' * 50)\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)\n",
    "\n",
    "with open(f\"{save_path}_full.json\", \"w\") as f:\n",
    "    json.dump(res_json_data, f)\n",
    "\n",
    "# Then you can run https://github.com/BradyFU/Video-MME/blob/main/evaluation/eval_your_results.py to get the score\n",
    "# python3 eval.py --results_file your_prediction_full.json --video_duration_type short,medium,long"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mistral",
   "language": "python",
   "name": "mistral"
  },
  "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}