Video-LLaVA-CinePile-Inference.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inference\n",
    "\n",
    "This collab is designed to run inference using Video-LLaVA fine tuned with CinePile pulling the model from Hugging Face Hub [mfarre/Video-LLaVA-7B-hf-CinePile](https://huggingface.co/mfarre/Video-LLaVA-7B-hf-CinePile).\n",
    "\n",
    "There are two ways to run it:\n",
    "* On any video - given a YouTube link and closed captions\n",
    "* On CinePile's test dataset\n",
    "\n",
    "Inference script based on [CinePile's Colab Notebook](https://colab.research.google.com/drive/1jDwvPoCsg9tck3dFhVCV-h3Ny6992wCr?usp=sharing)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "import shutil\n",
    "import PIL\n",
    "from PIL import Image\n",
    "import pathlib\n",
    "import matplotlib.pyplot as plt;\n",
    "import torch\n",
    "from transformers import VideoLlavaProcessor, VideoLlavaForConditionalGeneration, BitsAndBytesConfig\n",
    "from peft import PeftModel\n",
    "import subprocess\n",
    "import numpy as np\n",
    "from datasets import load_dataset\n",
    "\n",
    "from scenedetect import VideoManager\n",
    "from scenedetect import SceneManager\n",
    "from scenedetect.detectors import ContentDetector\n",
    "from scenedetect.scene_manager import save_images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "quantization_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_compute_dtype=torch.float16\n",
    ")\n",
    "\n",
    "model_id = \"LanguageBind/Video-LLaVA-7B-hf\"\n",
    "processor = VideoLlavaProcessor.from_pretrained(model_id)\n",
    "model = VideoLlavaForConditionalGeneration.from_pretrained(\n",
    "    model_id,\n",
    "    torch_dtype=torch.float16,\n",
    "    quantization_config=quantization_config,\n",
    "    device_map=\"auto\",\n",
    ")\n",
    "model = PeftModel.from_pretrained(model, \"mfarre/Video-LLaVA-7B-hf-CinePile\")\n",
    "\n",
    "MAX_GENERATED_LENGTH = 256 "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Auxiliar functions\n",
    "Mainly to handle video and extract inference frames from it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def download_video(video_url, filename, root):\n",
    "    \"\"\"\n",
    "    Download and convert a video from a URL and save it to a specified directory.\n",
    "\n",
    "    Parameters:\n",
    "    - video_url (str): The URL of the video to be downloaded.\n",
    "    - filename (str): The base name for the output file, without file extension.\n",
    "    - root (str): The root directory where the 'yt_videos' folder will be created.\n",
    "\n",
    "    Returns:\n",
    "    - tuple: A tuple containing the video URL and a boolean. The boolean is True if the\n",
    "      download and conversion was successful, and False otherwise.\n",
    "    \"\"\"\n",
    "\n",
    "    dir_path=f\"{root}/yt_videos\"\n",
    "\n",
    "    try:\n",
    "        vid_prefix = os.path.join(dir_path, filename)\n",
    "        full_command = [\n",
    "            \"yt-dlp\",\n",
    "            \"-S\",\n",
    "            \"height:224,ext:mp4:m4a\",\n",
    "            \"--recode\",\n",
    "            \"mp4\",\n",
    "            \"-o\",\n",
    "            f\"{vid_prefix}.mp4\",\n",
    "            video_url\n",
    "        ]\n",
    "\n",
    "        print(f'saving path: {vid_prefix}.mp4')\n",
    "\n",
    "        result = subprocess.run(full_command, capture_output=True, text=True)\n",
    "\n",
    "        if result.returncode == 0:\n",
    "            print(f\"Downloaded: {vid_prefix}; {video_url}\")\n",
    "            return video_url, True\n",
    "        else:\n",
    "            print(f\"Failed to download or convert {video_url}. Error: {result.stderr}\")\n",
    "            return video_url, False\n",
    "\n",
    "    except Exception as e:\n",
    "        print(f\"Exception during download or conversion of {video_url}: {e}\")\n",
    "        return video_url, False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_scenes(video_path, threshold=30.0):\n",
    "    \"\"\"\n",
    "    Detects important scenes in a video by analyzing changes between frames and identifying significant content changes that exceed a specified threshold.\n",
    "\n",
    "    Parameters:\n",
    "    video_path (str): The file path to the video file for which scenes are to be detected.\n",
    "    threshold (float): The sensitivity threshold for detecting scene changes.\n",
    "\n",
    "    Returns:\n",
    "    list of tuples: A list where each tuple contains the start and end `FrameTimecodes` of a detected scene.\n",
    "    \"\"\"\n",
    "\n",
    "    # Create a video manager object for the video.\n",
    "    video_manager = VideoManager([video_path])\n",
    "    scene_manager = SceneManager()\n",
    "\n",
    "    # Add ContentDetector algorithm (with a threshold).\n",
    "    scene_manager.add_detector(ContentDetector(threshold=threshold))\n",
    "\n",
    "    # Start the video manager and perform scene detection.\n",
    "    video_manager.set_downscale_factor()\n",
    "    video_manager.start()\n",
    "\n",
    "    # Perform scene detection and return scene list.\n",
    "    scene_manager.detect_scenes(frame_source=video_manager)\n",
    "\n",
    "    # Each scene is a tuple of (start, end) FrameTimecodes.\n",
    "    return scene_manager.get_scene_list()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_frames_from_scenes(video_path, scenes, output_folder):\n",
    "    \"\"\"\n",
    "    Extracts and saves the first frame from each detected scene in a video.\n",
    "\n",
    "    Parameters:\n",
    "    - video_path (str): The file path to the video from which frames are to be extracted.\n",
    "    - scenes (list): A list of scene boundaries or metadata that specifies where each scene begins and ends.\n",
    "    - output_folder (str): The directory path where the extracted frames should be saved.\n",
    "\n",
    "    Returns:\n",
    "    - None: The function saves the frames to the specified directory and does not return any value.\n",
    "    \"\"\"\n",
    "    # Ensure output directory exists.\n",
    "    if not os.path.exists(output_folder):\n",
    "        os.makedirs(output_folder)\n",
    "\n",
    "    # Initialize video manager for frame extraction.\n",
    "    video_manager = VideoManager([video_path])\n",
    "    video_manager.start()\n",
    "\n",
    "    # Save the first frame of each detected scene.\n",
    "    save_images(scenes, video_manager, num_images=1, output_dir=output_folder, image_name_template='$SCENE_NUMBER')\n",
    "\n",
    "    video_manager.release()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_uniform_frames(video_path, num_frames=10):\n",
    "    \"\"\"\n",
    "    This function takes a video file and returns a list of uniform frames from the video.\n",
    "    :param video_path: str, path to the video file\n",
    "    :param num_frames: int, number of uniform frames to return\n",
    "\n",
    "    :return: list of frames\n",
    "    \"\"\"\n",
    "    # check if path exists\n",
    "    if not os.path.exists(video_path):\n",
    "        raise FileNotFoundError(f\"Video file not found at {video_path}\")\n",
    "    cap = cv2.VideoCapture(video_path)\n",
    "    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n",
    "    uniform_frames = np.linspace(0, total_frames-1, num_frames, dtype=int)  # picking n frames uniformly from the video\n",
    "    # random_frames = random.sample(range(total_frames), num_frames)\n",
    "    frames = []\n",
    "    for frame_num in uniform_frames:\n",
    "        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_num)\n",
    "        ret, frame = cap.read()\n",
    "        if ret:\n",
    "            frames.append(frame)\n",
    "    return frames\n",
    "\n",
    "def save_frames_as_jpg(frames, frames_dir):\n",
    "    # create directory if it doesn't exist\n",
    "    pathlib.Path(frames_dir).mkdir(parents=True, exist_ok=True)\n",
    "    for i, frame in enumerate(frames):\n",
    "        cv2.imwrite(f\"{frames_dir}/{i}.jpg\", frame)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_video(test_dataset, root_dir, base_folder_name='new_yt_videos', max_num_frames=10, visualize=False, ques_idx=None):\n",
    "    if ques_idx is None:\n",
    "        ques_idx = random.randint(0, len(test_dataset))\n",
    "    print(ques_idx)\n",
    "    data = test_dataset[ques_idx]\n",
    "    clip_title, yt_link = data['yt_clip_title'], data['yt_clip_link']\n",
    "    print(clip_title, yt_link)\n",
    "\n",
    "    video_path = f\"{root_dir}/{base_folder_name}/{data['movie_name']}_{yt_link.split('/')[-1]}.mp4\"\n",
    "    frames_dir = f\"{root_dir}/{base_folder_name}_frames/_{data['movie_name']}_{yt_link.split('/')[-1]}\"\n",
    "    print(video_path)\n",
    "    if not os.path.exists(video_path):\n",
    "        download_video(yt_link, f\"{data['movie_name']}_{yt_link.split('/')[-1]}\", root=root_dir)\n",
    "        scenes = find_scenes(video_path)\n",
    "        save_frames_from_scenes(video_path, scenes, frames_dir)\n",
    "    else:\n",
    "        print(\"Skipping download, video exist\")\n",
    "\n",
    "\n",
    "    image_files = sorted([f for f in os.listdir(frames_dir) if f.endswith('.jpg')])[:-2]\n",
    "    if len(image_files) < max_num_frames:\n",
    "        shutil.rmtree(frames_dir)\n",
    "        frames = get_uniform_frames(video_path, max_num_frames+2)[:-2]\n",
    "        save_frames_as_jpg(frames, frames_dir)\n",
    "        image_files = sorted([f for f in os.listdir(frames_dir) if f.endswith('.jpg')])\n",
    "\n",
    "    img_file_paths = [os.path.join(frames_dir, image_file) for image_file in image_files]\n",
    "    num_frames = max_num_frames if len(img_file_paths) > max_num_frames else len(img_file_paths)\n",
    "    img_file_paths = [img_file_paths[i] for i in np.linspace(0, len(img_file_paths)-1, num_frames, dtype=int)]\n",
    "\n",
    "    if visualize:\n",
    "        num_cols = 5  # Number of columns in the grid\n",
    "        num_rows = (num_frames + num_cols - 1) // num_cols  # Calculate the necessary number of rows to display all images\n",
    "\n",
    "        fig, axs = plt.subplots(num_rows, num_cols, figsize=(15, 3 * num_rows))\n",
    "        axs = axs.flatten() if num_frames > 1 else [axs]\n",
    "\n",
    "        for i, img_path in enumerate(img_file_paths):\n",
    "            img = Image.open(img_path)\n",
    "            axs[i].imshow(img)\n",
    "            axs[i].set_title(os.path.basename(img_path), fontsize=8)\n",
    "            axs[i].axis('off')\n",
    "\n",
    "        for ax in axs[len(img_file_paths):]:  # Hide unused axes\n",
    "            ax.axis('off')\n",
    "\n",
    "        plt.tight_layout()\n",
    "        plt.show()\n",
    "\n",
    "    return data, img_file_paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_standalone_video(root_dir, yt_link, title, base_folder_name='new_yt_videos', max_num_frames=8, visualize=False):\n",
    "\n",
    "    video_path = f\"{root_dir}/{base_folder_name}/{title}_{yt_link.split('/')[-1]}.mp4\"\n",
    "    frames_dir = f\"{root_dir}/{base_folder_name}_frames/_{title}_{yt_link.split('/')[-1]}\"\n",
    "    print(video_path)\n",
    "    if not os.path.exists(video_path):\n",
    "        download_video(yt_link, f\"{title}_{yt_link.split('/')[-1]}\", root=root_dir)\n",
    "        scenes = find_scenes(video_path)\n",
    "        save_frames_from_scenes(video_path, scenes, frames_dir)\n",
    "    else:\n",
    "        print(\"Skipping download, video exist\")\n",
    "\n",
    "\n",
    "    image_files = sorted([f for f in os.listdir(frames_dir) if f.endswith('.jpg')])[:-2]\n",
    "    if len(image_files) < max_num_frames:\n",
    "        shutil.rmtree(frames_dir)\n",
    "        frames = get_uniform_frames(video_path, max_num_frames+2)[:-2]\n",
    "        save_frames_as_jpg(frames, frames_dir)\n",
    "        image_files = sorted([f for f in os.listdir(frames_dir) if f.endswith('.jpg')])\n",
    "\n",
    "    img_file_paths = [os.path.join(frames_dir, image_file) for image_file in image_files]\n",
    "    num_frames = max_num_frames if len(img_file_paths) > max_num_frames else len(img_file_paths)\n",
    "    img_file_paths = [img_file_paths[i] for i in np.linspace(0, len(img_file_paths)-1, num_frames, dtype=int)]\n",
    "\n",
    "    if visualize:\n",
    "        num_cols = 5  # Number of columns in the grid\n",
    "        num_rows = (num_frames + num_cols - 1) // num_cols  # Calculate the necessary number of rows to display all images\n",
    "\n",
    "        fig, axs = plt.subplots(num_rows, num_cols, figsize=(15, 3 * num_rows))\n",
    "        axs = axs.flatten() if num_frames > 1 else [axs]\n",
    "\n",
    "        for i, img_path in enumerate(img_file_paths):\n",
    "            img = Image.open(img_path)\n",
    "            axs[i].imshow(img)\n",
    "            axs[i].set_title(os.path.basename(img_path), fontsize=8)\n",
    "            axs[i].axis('off')\n",
    "\n",
    "        for ax in axs[len(img_file_paths):]:  # Hide unused axes\n",
    "            ax.axis('off')\n",
    "\n",
    "        plt.tight_layout()\n",
    "        plt.show()\n",
    "\n",
    "    return img_file_paths"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inference on wild videos\n",
    "Run this section for inference on a YouTube video of your choice\n",
    "Note that the answers of the model have a strong bias towards the CinePile format. We clean it up to keep only the first sentence as it is usually where the value is."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ROOT_DIR = 'cinepileevaldata'\n",
    "video_link = \"https://www.youtube.com/watch?v=J1NThLImq9I\"\n",
    "own_title = \"carretera-perdida\"\n",
    "img_file_paths = process_standalone_video(ROOT_DIR, video_link, own_title , base_folder_name='yt_videos', max_num_frames=8, visualize=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Hand-made prompt. Note three important aspects: \n",
    "- It starts with `USER:` and ends with `ASSISTANT:`\n",
    "- The `<video>` tag is necessary to insert the tokens from the video frames. It is placed between `USER:` and `ASSISTANT:`\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"\"\"\n",
    "\"USER: <prompt>You will be provided with subtitles from a specific scene of a movie and a few frames from that scene. After going through the movie scene and seeing the frames, please answer the question that follows.\n",
    "\n",
    "**Subtitles:** \n",
    "<subtitle> Súbete a mi lado que yo conduzco hoy\n",
    "<subtitle> Voy a ir sin manos y tú sin cinturón\n",
    "<subtitle> Pon una canción que no hable de amor\n",
    "<subtitle> Dale a toda hostia que reviente el altavoz, y\n",
    "<subtitle> Acelero a fondo en contra dirección\n",
    "<subtitle> El que se aparte antes es el perdedor\n",
    "<subtitle> Vamos a llevar al límite el motor\n",
    "<subtitle> Apaga la colilla con la punta del tacón\n",
    "<subtitle> La brisa en la cara\n",
    "<subtitle> La luz de la noche\n",
    "<subtitle> La ventanilla bajada\n",
    "<subtitle> Me mira y se coloca el escote\n",
    "<subtitle> La brisa en la cara\n",
    "<subtitle> La luz de la noche\n",
    "<subtitle> La ventanilla bajada\n",
    "<subtitle> La música sonando\n",
    "<subtitle> Y el corazón, a ciento veinte pulsaciones\n",
    "<subtitle> Saltan flashes por exceso de velocidad\n",
    "<subtitle> Rayas con la uña el vaho que hay en el cristal\n",
    "<subtitle> Los pies en la guantera y la cabeza para atrás\n",
    "<subtitle> Se baja la cremallera hasta la mitad\n",
    "<subtitle> Una curva cerrada voy rozando el quitamiedos\n",
    "<subtitle> Cojo la salida que nos lleva hacia un sendero\n",
    "<subtitle> El aullido de una loba se escucha a lo lejos\n",
    "<subtitle> Me ha clava'o los dos colmillos en el cuello\n",
    "<subtitle> La brisa en la cara\n",
    "<subtitle> La luz de la noche\n",
    "<subtitle> La ventanilla bajada\n",
    "<subtitle> Me mira y se coloca el escote\n",
    "<subtitle> La brisa en la cara\n",
    "<subtitle> La luz de la noche\n",
    "<subtitle> La ventanilla bajada\n",
    "<subtitle> La música sonando\n",
    "<subtitle> Y el corazón, a ciento veinte pulsaciones\n",
    "<subtitle> La brisa\n",
    "<subtitle> En la cara\n",
    "<subtitle> La luz de la noche\n",
    "<subtitle> La luz de la noche (ahh)\n",
    "<subtitle> La brisa en la cara\n",
    "<subtitle> La luz de la noche (la luz de la noche)\n",
    "<subtitle> La ventanilla bajada (ah)\n",
    "<subtitle> Me mira y se coloca el escote (u-u-uh)\n",
    "<subtitle> La brisa en la cara (la brisa)\n",
    "<subtitle> La luz de la noche (de la noche)\n",
    "<subtitle> La ventanilla bajada\n",
    "<subtitle> La música sonando\n",
    "<subtitle> Y el corazón, a ciento veinte pulsaciones\n",
    "<video>\n",
    "Question: {question}\n",
    "\n",
    "ASSISTANT:\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "formatted_prompt = prompt.format(question=\"What can you say about the driving style?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Loading the video and processing it together with the prompt\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clip = np.stack([np.array((PIL.Image.open(f)).convert('RGB')) for f in img_file_paths])\n",
    "inputs = processor(text=formatted_prompt, videos=clip, return_tensors=\"pt\").to(model.device)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate token IDs and decode back into text\n",
    "generated_ids = model.generate(**inputs, max_new_tokens=MAX_GENERATED_LENGTH)\n",
    "generated_texts = processor.batch_decode(generated_ids, skip_special_tokens=True)\n",
    "\n",
    "# As the model is fine-tuned for CinePile, it might output too much noise on other tasks. Here we do a quick clean-up of the output.\n",
    "print(generated_texts[0].split(\"ASSISTANT:\")[1].split(\"\\n\")[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "formatted_prompt = prompt.format(question=\"How many 🚬 did he smoke?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clip = np.stack([np.array((PIL.Image.open(f)).convert('RGB')) for f in img_file_paths])\n",
    "inputs = processor(text=formatted_prompt, videos=clip, return_tensors=\"pt\").to(model.device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate token IDs and decode back into text\n",
    "generated_ids = model.generate(**inputs, max_new_tokens=MAX_GENERATED_LENGTH)\n",
    "generated_texts = processor.batch_decode(generated_ids, skip_special_tokens=True)\n",
    "\n",
    "# As the model is fine-tuned for CinePile, it might output too much noise on other tasks. Here we do a quick clean-up of the output.\n",
    "print(generated_texts[0].split(\"ASSISTANT:\")[1].split(\"\\n\")[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CinePile Evaluation\n",
    "Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ROOT_DIR = 'cinepileevaldata'\n",
    "MAX_LENGTH = 256\n",
    "MAX_NUM_FRAMES = 10\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Prompt formatting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ans_key_map = {0: 'A', 1: 'B', 2: 'C', 3: 'D', 4: 'E'}\n",
    "vision_and_language_dependence_prompt = '''USER: <prompt>You will be provided with subtitles from a specific scene of a movie and a few frames from that scene. After going through the movie scene and seeing the frames, please answer the question that follows. The question will have five possible answers labeled A, B, C, D, and E, please try to provide the most probable answer in your opinion. Your output should be just one of A,B,C,D,E and nothing else.\n",
    "\n",
    "**Output Format:**\n",
    "    **Answer:** <Option_key>\n",
    "\n",
    "**Subtitles:** \\n{subs}\\n<video>\\nQuestion: {question}\n",
    "\n",
    "Note: Follow the output format strictly. Only answer with the option key (A, B, C, D, E) and nothing else.\n",
    "ASSISTANT:'''\n",
    "\n",
    "def format_question_and_options(question, options):\n",
    "    \"\"\"\n",
    "    Formats a question and a list of options into a single string with options labeled A, B, C, etc.\n",
    "\n",
    "    Parameters:\n",
    "    - question (str): The question to be formatted.\n",
    "    - options (list of str): The options for the question.\n",
    "\n",
    "    Returns:\n",
    "    - str: The formatted question and options.\n",
    "    \"\"\"\n",
    "    formatted_string = f\"{question}\\n\"\n",
    "    option_labels = [chr(ord('A') + i) for i in range(len(options))]  # Generate option labels dynamically\n",
    "\n",
    "    for label, option in zip(option_labels, options):\n",
    "        formatted_string += f\"- {label}) {option}\\n\"\n",
    "\n",
    "    return formatted_string\n",
    "\n",
    "def get_prompt(data):\n",
    "    formatted_subs = data['subtitles']\n",
    "    options = data['choices']\n",
    "    formatted_question = format_question_and_options(data['question'], options)\n",
    "\n",
    "    prompt = vision_and_language_dependence_prompt.format(subs=formatted_subs, question=formatted_question)\n",
    "    return prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cinepile = load_dataset(\"tomg-group-umd/cinepile\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We ran inference through all the clips"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = []\n",
    "for q in range(0,len(cinepile['test'])):\n",
    "    data, img_file_paths = process_video(cinepile['test'], ROOT_DIR, base_folder_name='yt_videos', max_num_frames=MAX_NUM_FRAMES, visualize=False, ques_idx=q)\n",
    "    prompt = get_prompt(data)\n",
    "    clip = np.stack([np.array((PIL.Image.open(f)).convert('RGB')) for f in img_file_paths])\n",
    "    inputs = processor(text=prompt, videos=clip, return_tensors=\"pt\").to(model.device)\n",
    "    generated_ids = model.generate(**inputs, max_new_tokens=MAX_LENGTH)\n",
    "    generated_texts = processor.batch_decode(generated_ids, skip_special_tokens=True)\n",
    "    #answer = generated_texts[0][-1]\n",
    "    answer = generated_texts[0].split(\"ASSISTANT:\")[1][0]\n",
    "\n",
    "    ans_key_opt = ans_key_map[data['answer_key_position']]\n",
    "    correct = False\n",
    "    if ans_key_opt == answer:\n",
    "        correct = True\n",
    "    results.append({\"ques_idx\":q,\"full_answer\":generated_texts[0],\"parsed_answer\":answer,\"correct\":correct})\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We present the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Compiling questions per category\n",
    "categories = dict()\n",
    "total_questions = 0\n",
    "total_correct = 0\n",
    "\n",
    "for elem in cinepile['test']:\n",
    "    if elem['question_category'] not in categories:\n",
    "        categories[elem['question_category']] = 1\n",
    "    else:\n",
    "        categories[elem['question_category']] += 1\n",
    "\n",
    "\n",
    "\n",
    "benchmark = {key: 0 for key in categories}\n",
    "benchmark_visual = {key: 0 for key in categories}\n",
    "\n",
    "for idx, answer in enumerate(results):\n",
    "    total_questions +=1\n",
    "    if answer['correct']:\n",
    "        benchmark[cinepile['test'][idx]['question_category']] += 1\n",
    "        total_correct += 1\n",
    "        if cinepile['test'][idx]['visual_reliance'] == \"1\":\n",
    "            benchmark_visual[cinepile['test'][idx]['question_category']] += 1\n",
    "            \n",
    "\n",
    "performance = {key: 100.0 * benchmark[key] / categories[key] for key in categories}\n",
    "\n",
    "\n",
    "print(f\"Average score: {100.0*total_correct/total_questions}\")\n",
    "print(f\"Per category score: {performance}\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}