{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Demo: Qwen3-VL-8B answering one CHNB question with 64 uniform frames\n",
    "\n",
    "Single-question end-to-end run on a local GPU. Picks the first benchmark question whose mp4 is on disk, samples 64 uniform frames, sends them to Qwen3-VL-8B-Instruct via `transformers`, and prints the answer alongside the gold letter.\n",
    "\n",
    "```bash\n",
    "pip install -r requirements.txt\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "from pathlib import Path\n",
    "\n",
    "_PROJECT = Path.cwd().resolve()\n",
    "while not (_PROJECT / 'util').is_dir() and _PROJECT != _PROJECT.parent:\n",
    "    _PROJECT = _PROJECT.parent\n",
    "sys.path.insert(0, str(_PROJECT))\n",
    "print(f'Project root: {_PROJECT}')\n",
    "\n",
    "from util.benchmark import load_benchmark\n",
    "from util.frames import extract_frames_uniform\n",
    "from util.parsers import parse_mcq_json\n",
    "from util.paths import BENCHMARK_PATH, VIDEO_DIR\n",
    "from util.prompts import HEADER_VIDEO_FRAMES, mcq_prompt_json\n",
    "\n",
    "df = load_benchmark(BENCHMARK_PATH)\n",
    "print(f'{len(df)} questions across {df[\"video_file_id\"].nunique()} movies')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pick the first question whose mp4 is on disk.\n",
    "available = [v for v in df['video_file_id'].drop_duplicates() if (VIDEO_DIR / f'{v}.mp4').exists()]\n",
    "if not available:\n",
    "    raise FileNotFoundError(\n",
    "        f'No mp4s found under {VIDEO_DIR}. Run `bash scripts/download_movies.sh` first '\n",
    "        f'(or place a single mp4 there to test).'\n",
    "    )\n",
    "vid = available[0]\n",
    "row = df[df['video_file_id'] == vid].iloc[0].to_dict()\n",
    "print(f'video_file_id : {vid}')\n",
    "print(f'category      : {row[\"category\"]}')\n",
    "print(f'question      : {row[\"question\"]}')\n",
    "for L in 'ABCD':\n",
    "    print(f'  ({L}) {row[f\"option_{L}\"]}')\n",
    "print(f'gold          : {row[\"gold\"]}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sample 64 uniform frames from the source mp4.\n",
    "NUM_FRAMES = 64\n",
    "video_path = VIDEO_DIR / f'{vid}.mp4'\n",
    "frames, timestamps = extract_frames_uniform(video_path, NUM_FRAMES, low_memory=False)\n",
    "print(f'Extracted {len(frames)} frames spanning {timestamps[0]:.1f}s — {timestamps[-1]:.1f}s')\n",
    "frames[0]   # display the first frame inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load Qwen3-VL-8B-Instruct. First run will download ~17 GB from Hugging Face.\n",
    "import torch\n",
    "from transformers import AutoModelForImageTextToText, AutoProcessor\n",
    "from qwen_vl_utils import process_vision_info\n",
    "\n",
    "MODEL_ID = 'Qwen/Qwen3-VL-8B-Instruct'\n",
    "\n",
    "processor = AutoProcessor.from_pretrained(\n",
    "    MODEL_ID, trust_remote_code=True,\n",
    "    min_pixels=256 * 28 * 28, max_pixels=512 * 28 * 28,\n",
    ")\n",
    "model = AutoModelForImageTextToText.from_pretrained(\n",
    "    MODEL_ID, dtype=torch.bfloat16, device_map='auto',\n",
    "    attn_implementation='flash_attention_2', trust_remote_code=True,\n",
    ")\n",
    "print(f'Loaded {MODEL_ID} on {next(model.parameters()).device}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build the chat-template input and run inference.\n",
    "header = HEADER_VIDEO_FRAMES.format(n_frames=NUM_FRAMES)\n",
    "prompt = mcq_prompt_json(row, header=header)\n",
    "print('--- PROMPT ---')\n",
    "print(prompt[:500] + ('...' if len(prompt) > 500 else ''))\n",
    "\n",
    "content = [{'type': 'image', 'image': img} for img in frames]\n",
    "content.append({'type': 'text', 'text': prompt})\n",
    "messages = [\n",
    "    {'role': 'system', 'content': 'You answer movie MCQs. Output only valid JSON.'},\n",
    "    {'role': 'user', 'content': content},\n",
    "]\n",
    "text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
    "image_inputs, video_inputs = process_vision_info(messages)\n",
    "inputs = processor(text=[text], images=image_inputs, videos=video_inputs,\n",
    "                   return_tensors='pt', padding=True)\n",
    "inputs = {k: (v.to(model.device) if hasattr(v, 'to') else v) for k, v in inputs.items()}\n",
    "\n",
    "with torch.no_grad():\n",
    "    gen_ids = model.generate(**inputs, max_new_tokens=2048, temperature=0.1, do_sample=True)\n",
    "in_len = inputs['input_ids'].shape[1]\n",
    "raw = processor.batch_decode(gen_ids[:, in_len:], skip_special_tokens=True)[0].strip()\n",
    "print('--- RAW RESPONSE ---')\n",
    "print(raw)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Parse the answer.\n",
    "pred = parse_mcq_json(raw, qid=row['qid'])\n",
    "print(f'Gold: {row[\"gold\"]}')\n",
    "print(f'Pred: {pred}')\n",
    "print('Correct' if pred == row['gold'] else 'Wrong')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}