{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d762edec-c0d4-44b4-baae-127d1bab372e",
   "metadata": {},
   "source": [
    "## Gemini demo\n",
    "\n",
    "This notebook assumes you:\n",
    "\n",
    "- ran the script in `process_tvqa` -- you need the reconstructed clips and candidate lists (a.k.a. multiple_choice related files)\n",
    "- have access to a Vertex AI Workbench instance (where this notebook was tested)\n",
    "- have enabled the Gemini API\n",
    "- store your files in a Cloud Storage bucket in Google Cloud Platform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb368613-0421-4673-834b-eb1aa57eeb1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import re\n",
    "import json\n",
    "import datetime\n",
    "from google import genai\n",
    "from pydantic import BaseModel\n",
    "from tqdm.notebook import tqdm\n",
    "from google.cloud import storage\n",
    "from IPython.display import HTML, Markdown, display\n",
    "from google.genai.types import GenerateContentConfig, Part"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0406d3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def list_files_in_folder(bucket_name, folder_path):\n",
    "    \"\"\"Lists all files in a folder within a Google Cloud Storage bucket.\"\"\"\n",
    "\n",
    "    storage_client = storage.Client()\n",
    "    bucket = storage_client.bucket(bucket_name)\n",
    "    blobs = bucket.list_blobs(prefix=folder_path)  # Use prefix to filter\n",
    "\n",
    "    file_list = []  # Initialize an empty list\n",
    "    for blob in blobs:\n",
    "        if not blob.name.endswith('/'):  # Exclude \"folders\" (objects ending in /)\n",
    "            file_list.append(blob.name)\n",
    "\n",
    "    return file_list\n",
    "\n",
    "def read_json_from_gcs(bucket_name, blob_name):\n",
    "    \"\"\"Reads a JSON file from Google Cloud Storage.\"\"\"\n",
    "\n",
    "    try:\n",
    "        storage_client = storage.Client()\n",
    "        bucket = storage_client.bucket(bucket_name)\n",
    "        blob = bucket.blob(blob_name)\n",
    "\n",
    "        # Download the file content as bytes\n",
    "        file_content = blob.download_as_bytes()\n",
    "\n",
    "        # Decode the bytes and parse the JSON\n",
    "        data = json.loads(file_content)\n",
    "\n",
    "        return data\n",
    "\n",
    "    except Exception as e:\n",
    "        print(f\"Error reading file: {e}\")\n",
    "        return None\n",
    "\n",
    "\n",
    "def write_json_to_gcs(bucket_name, blob_name, data):\n",
    "    \"\"\"Writes JSON data to a Google Cloud Storage blob.\"\"\"\n",
    "    try:\n",
    "        storage_client = storage.Client()\n",
    "        bucket = storage_client.bucket(bucket_name)\n",
    "        blob = bucket.blob(blob_name)\n",
    "        blob.upload_from_string(data, content_type='application/json')\n",
    "    except Exception as e:\n",
    "        print(f\"Error writing JSON to GCS: {e}\")\n",
    "\n",
    "def blob_exists(bucket_name, blob_name):\n",
    "    \"\"\"Checks if a blob exists in a Google Cloud Storage bucket.\"\"\"\n",
    "\n",
    "    try:\n",
    "        storage_client = storage.Client()\n",
    "        bucket = storage_client.bucket(bucket_name)\n",
    "        blob = bucket.blob(blob_name)\n",
    "        return blob.exists()\n",
    "\n",
    "    except Exception as e:\n",
    "        print(f\"Error checking blob existence: {e}\")\n",
    "        return False\n",
    "\n",
    "def extract_clip_id(blob_name):\n",
    "    \"\"\"Extracts the clip ID from a blob name.\"\"\"\n",
    "\n",
    "    try:\n",
    "        # Updated regex to handle different show names\n",
    "        match = re.search(r\"[a-z]+_s\\d+e\\d+_seg\\d+_clip_\\d+\", blob_name)\n",
    "        if match:\n",
    "            return match.group(0)\n",
    "        else:\n",
    "            return None  # Clip ID not found\n",
    "\n",
    "    except Exception as e:\n",
    "        print(f\"Error extracting clip ID: {e}\")\n",
    "        return None        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6827c2e7-335c-4c09-810a-9f71fa1f9daf",
   "metadata": {},
   "outputs": [],
   "source": [
    "PROJECT_ID = \"YOUR_PROJECT_ID\" \n",
    "LOCATION = \"YOUR_LOCATION\"\n",
    "\n",
    "client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "717badb0-d745-43c8-a433-7548ad3b8701",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96c47824-bad4-4389-980e-f71fa59d1160",
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_ID = \"gemini-2.0-flash\"  \n",
    "\n",
    "bucket_name = \"YOUR_BUCKET_NAME\"\n",
    "folder_path = \"YOUR_DATA_FOLDER/\" \n",
    "\n",
    "files = list_files_in_folder(bucket_name, folder_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25ed7cfb-906d-49c3-aa35-a5874973017a",
   "metadata": {},
   "outputs": [],
   "source": [
    "file_dict = {}\n",
    "\n",
    "for file in files:\n",
    "    try:\n",
    "        show_id = file.split('/')[1]\n",
    "        clip_id = file.split('/')[2].split('.')[0]\n",
    "    except:\n",
    "        print(file)\n",
    "  \n",
    "    if '.DS_Store' in str(file) or '.ipynb_checkpoints' in str(file):\n",
    "        continue\n",
    "        \n",
    "    if '.annotated_video.mp4' in file:\n",
    "        file_dict.setdefault(show_id, {}).setdefault(clip_id, {}).update(\n",
    "            {'video_path': file}\n",
    "        ) \n",
    "    elif '.audio.mp3' in file:\n",
    "        file_dict.setdefault(show_id, {}).setdefault(clip_id, {}).update(\n",
    "            {'audio_path': file}\n",
    "        )        \n",
    "    elif '.multiple_choice.json' in file:\n",
    "        file_dict.setdefault(show_id, {}).setdefault(clip_id, {}).update(\n",
    "            {'multiple_choice_path': file}\n",
    "        )   \n",
    "    elif '.multiple_choice/' in file:\n",
    "        file_dict.setdefault(show_id, {})\\\n",
    "                 .setdefault(clip_id, {})\\\n",
    "                 .setdefault('multiple_choice_image_paths', [])\\\n",
    "                 .append(file)\n",
    "    elif '.frames/' in file:\n",
    "        file_dict.setdefault(show_id, {})\\\n",
    "                 .setdefault(clip_id, {})\\\n",
    "                 .setdefault('frame_paths', [])\\\n",
    "                 .append(file)        \n",
    "    elif '.subtitles.json' in file:\n",
    "        file_dict.setdefault(show_id, {}).setdefault(clip_id, {}).update(\n",
    "            {'subtitle_path': file}\n",
    "        )  \n",
    "     \n",
    "    else:\n",
    "        print(file)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b604fb83-c74c-4b12-a98f-923875044e43",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConversationalRoles(BaseModel):\n",
    "    line_index: int\n",
    "    reply_to: int\n",
    "    speaker: str\n",
    "    addressees: list[str]\n",
    "    side_participants: list[str]\n",
    "\n",
    "class ClipRoles(BaseModel):\n",
    "    clip_roles: list[ConversationalRoles]    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77fc816b-51e1-42c6-879f-ce9ff904307a",
   "metadata": {},
   "outputs": [],
   "source": [
    "system_instruction = \"\"\"You are a video analysis assistant.  Your task is to analyze the conversations in a video clip and its associated subtitles. For each dialogue line, you will:\n",
    "\n",
    "*   determine what previous line it is replying to\n",
    "*   determine the speaker, addressees, and side-participants\n",
    "\n",
    "Here's how to determine the reply-to relationship between utterances to resolve conversational threads:\n",
    "\n",
    "*   The reply-to structure gives us information about floor-claiming and topical change within the clip.\n",
    "*   The character is saying this line because they want to respond to that previous line. What previous line is this current line replying to?\n",
    "*   If the speaker of the last line is the same, you can treat it as continuation and put the index of last line as the reply-to.\n",
    "*   When there is a noticeable change in topic and distribution of other participants' attention, and no previous line **triggers** this current line: write the current line index, indicating the current line replies to itself.\n",
    "\n",
    "Here's how to determine each role:\n",
    "\n",
    "*   **Speaker:**  The character who is speaking the line.  Infer this from lip movements, body language, and the context of the dialogue. If a character finishes one line and immediately starts another (very short pause), assume it's the same speaker, UNLESS there's a clear visual indication of a scene or speaker change (e.g., a camera cut to a different person starting to speak).\n",
    "*   **Addressee(s):** The character(s) the speaker is *directly* addressing. Use these cues:\n",
    "    *   **Eye Contact:** The most important cue. Who is the speaker looking at?\n",
    "    *   **Body Orientation:** Is the speaker's body turned towards a particular person or group?\n",
    "    *   **Dialogue Context:** Does the line contain a name, pronoun (\"you\"), or clearly refer to a specific individual or group?  (\"Hey, John...\" or \"You all need to...\")\n",
    "    *   **Reactions:** If a character reacts immediately and strongly to a line (e.g., nods, responds verbally, shows surprise), they are likely an addressee.\n",
    "    *   If the speaker seems to be talking to everyone present, list all characters who appear to be paying attention.\n",
    "    *   If the speaker is talking to a crowd of unidentifiable characters, write \"crowd\".\n",
    "    *   If the speaker is talking to themselves, or no one in particular, write \"none\".\n",
    "*   **Side-Participant(s):**  Any character(s) visible in the scene *during the line's timeframe* who are *not* the speaker or addressees. They are present, and their presence is known to other participants. They can potentially join the conversation at any time.\n",
    "    *   If it is not possible to confidently determine if someone is a side-participant, write \"unknown\".\n",
    "    *   If there are no side-participants, write \"none\".\n",
    "\n",
    "\n",
    "**Input:**\n",
    "\n",
    "You will receive a list of subtitle entries.  Each entry will be a dictionary with the following keys:\n",
    "*   `\"line_index\"`: (int) The index of the current entry (subtitle line).\n",
    "*   `\"start_time\"`: (float) The start time of the subtitle line in seconds.\n",
    "*   `\"end_time\"`: (float) The end time of the subtitle line in seconds.\n",
    "*   `\"text\"`: (string) The text of the dialogue line.\n",
    "\n",
    "You will also receive a list of potential participants for you to assign roles from. You must pick from this list.\n",
    "\n",
    "With all this information, analyze the video segment corresponding to the `start_time` and `end_time` of each subtitle entry.\n",
    "\n",
    "**Output:**\n",
    "\n",
    "Provide your output in JSON format, mirroring the structure of the input.  For *each* subtitle entry, add the following keys:\n",
    "*   `\"line_index\"`: (int) The line being analyzed.\n",
    "*   `\"reply_to\"`: (int) The line index that this current line replies to, could be the same as the current line index or any previous line index.\n",
    "*   `\"speaker\"`: (string) The name of the speaker. If you cannot determine the speaker, use \"unknown\".\n",
    "*   `\"addressees\"`: (list of strings) A list of the names of the addressee(s).  This can be an empty list (`[]`) if there are no direct addressees, or `[\"none\"]` if the speaker is speaking generally but to no one in particular.\n",
    "*   `\"side_participants\"`: (list of strings) A list of the names of the side-participant(s). This can be an empty list (`[]`), `[\"none\"]`, or `[\"unknown\"]`.\n",
    "\n",
    "This corresponds to the data model format:\n",
    "\n",
    "```python\n",
    "class ConversationalRoles(BaseModel):\n",
    "    reply_to: int  # Index of the line being replied to\n",
    "    speaker: str   # Speaker of the line, or \"unknown\"\n",
    "    addressees: list[str]  # List of names, [\"crowd\"], [\"none\"], or [\"unknown\"]\n",
    "    side_participants: list[str] # List of names, [\"crowd\"], [\"none\"], or [\"unknown\"]\n",
    "    \n",
    "class ClipRoles(BaseModel):\n",
    "    clip_roles: list[ConversationalRoles]        \n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efd719cd-a93a-41f8-ac93-2c504ea33ff7",
   "metadata": {},
   "outputs": [],
   "source": [
    "prefix = 'PRED_gemini'\n",
    "\n",
    "for show_id in tqdm(file_dict):\n",
    "    for clip_id, path_dict in tqdm(file_dict[show_id].items(), total=len(file_dict[show_id])):\n",
    "        pred_blob_name = f\"{prefix}/{show_id}/{clip_id}.{MODEL_ID}.json\"\n",
    "        \n",
    "        if blob_exists(bucket_name, pred_blob_name):\n",
    "            continue\n",
    "\n",
    "        prompt = \"\"\"Analyze this video. Pay attention to the bounding boxes and character captions.\n",
    "        \n",
    "Subtitles:\n",
    "\"\"\"            \n",
    "        subtitle_blob_name = path_dict['subtitle_path']\n",
    "        subtitle_data = read_json_from_gcs(bucket_name, subtitle_blob_name)\n",
    "        subtitle_list = [{\"line_index\": line_idx, \"start_time\": item[0], \"end_time\": item[1], \"text\": item[2]} for line_idx, item in enumerate(subtitle_data)]\n",
    "        for subtitle in subtitle_list:\n",
    "            prompt += json.dumps(subtitle) + '\\n'        \n",
    "\n",
    "        multiple_choice_blob_name = path_dict['multiple_choice_path']\n",
    "        multiple_choice_data = read_json_from_gcs(bucket_name, multiple_choice_blob_name)\n",
    "        \n",
    "        prompt += f\"\\n Select only from the following participants:\"\n",
    "        \n",
    "        mc_folder_prefix = f\"{show_id}/{clip_id}.multiple_choice/\"\n",
    "        mc_images = path_dict['multiple_choice_image_paths']\n",
    "\n",
    "        candidate_parts = []\n",
    "        for participant in multiple_choice_data:\n",
    "            candidate_img = None\n",
    "            participant_safe_char = participant.replace(\" \", \"_\")\n",
    "            # Look for a file whose name contains _{participant_safe_char}.jpg.\n",
    "            for img_filename in mc_images:\n",
    "                if f\"_{participant_safe_char}.jpg\" in img_filename:\n",
    "                    candidate_img = f\"gs://{bucket_name}/{img_filename}\"\n",
    "                    break\n",
    "            # Only add an image part if a candidate image was found.\n",
    "            candidate_parts.append(participant)            \n",
    "            if candidate_img:\n",
    "                candidate_parts.append(\n",
    "                    Part.from_uri(file_uri=candidate_img, mime_type=\"image/jpeg\")\n",
    "                )     \n",
    "\n",
    "        try:\n",
    "            video_part = Part.from_uri(\n",
    "                file_uri= f\"gs://{bucket_name}/\" + path_dict['video_path'],\n",
    "                mime_type=\"video/mp4\",\n",
    "            )     \n",
    "            \n",
    "            response = client.models.generate_content(\n",
    "                model=MODEL_ID,\n",
    "                contents=[\n",
    "                    prompt,                \n",
    "                    candidate_parts,\n",
    "                    video_part,\n",
    "                ],\n",
    "                config=GenerateContentConfig(\n",
    "                    system_instruction=system_instruction,\n",
    "                    response_mime_type=\"application/json\",\n",
    "                    response_schema=ClipRoles,        \n",
    "                ),\n",
    "            )       \n",
    "            pred = response.text\n",
    "            write_json_to_gcs(bucket_name, pred_blob_name, pred)\n",
    "        except Exception as e:\n",
    "            print(e)\n",
    "            continue        "
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "conda-base-py",
   "name": "workbench-notebooks.m128",
   "type": "gcloud",
   "uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/workbench-notebooks:m128"
  },
  "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}