{ "nbformat": 4, "nbformat_minor": 0, "metadata": { "colab": { "provenance": [], "machine_shape": "hm", "gpuType": "L4" }, "kernelspec": { "name": "python3", "display_name": "Python 3" }, "language_info": { "name": "python" }, "accelerator": "GPU" }, "cells": [ { "cell_type": "markdown", "source": [ "# Demo: Post-process TVQA for TV-MMPC\n", "\n", "This Colab notebook demonstrates how to post-process **TVQA** data for use with **TV-MMPC**. It assumes you already have access to:\n", "\n", "- The **original TVQA files** (please see the [original project](https://nlp.cs.unc.edu/data/jielei/tvqa/tvqa_public_html/download_tvqa.html) page -- note it's TVQA not TVQA+)\n", "- The **annotation files** available [here](https://doi.org/10.7910/DVN/4KUKUL).\n", "- A HuggingFace token for WhisperX access.\n", "\n", "## Required Files\n", "\n", "For this demonstration, please make sure the following files are available for each clip: for example, `s01e02_seg02_clip_01` points to the *The Big Bang Theory* season 1, episode 1, segment 2, clip 1 from TVQA (we retain the original file names):\n", "\n", "- `s01e02_seg02_clip_01.mp3` -- extracted audio from the TVQA clip\n", "- `s01e02_seg02_clip_01.srt` -- subtitles for the clip in `.srt` format\n", "- `s01e02_seg02_clip_01.annotation.json` -- annotations for this clip\n", "\n", "The notebook was tested in a Colab environment with an L4 GPU with the files listed above." ], "metadata": { "id": "l-bIFpWKMSTM" } }, { "cell_type": "markdown", "source": [ "## Setup for Colab" ], "metadata": { "id": "lSP1obvtHxqV" } }, { "cell_type": "code", "source": [ "!pip install pysrt pandas rapidfuzz scikit-learn" ], "metadata": { "id": "5jbOrtA3T8W9" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "source": [ "N.B. It takes some clunky workaround to make WhisperX work on Colab (this will take around 5 minutes):" ], "metadata": { "id": "jq6KhQQ9UjTW" } }, { "cell_type": "code", "source": [ "!pip uninstall torch torchvision torchaudio -y\n", "\n", "# Workaround from: https://github.com/m-bain/whisperX/issues/1027#issuecomment-2627525081\n", "!pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cu121\n", "\n", "# WhisperX-related packages:\n", "!pip install ctranslate2==4.4.0\n", "!pip install faster-whisper==1.1.0\n", "# !pip install git+https://github.com/m-bain/whisperx.git\n", "!pip install whisperx==3.3.1\n", "\n", "!apt-get update\n", "!apt-get install libcudnn8=8.9.2.26-1+cuda12.1\n", "!apt-get install libcudnn8-dev=8.9.2.26-1+cuda12.1\n", "\n", "!python -c \"import torch; torch.backends.cuda.matmul.allow_tf32 = True; torch.backends.cudnn.allow_tf32 = True\"" ], "metadata": { "id": "hCZ6dQFTHpeR" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "source": [ "To authenticate for WhisperX:" ], "metadata": { "id": "GWdUWX_CUrP1" } }, { "cell_type": "code", "source": [ "from huggingface_hub import notebook_login\n", "notebook_login()" ], "metadata": { "id": "75bh7w7NT1hV" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "source": [ "## Main" ], "metadata": { "id": "8BBkDnY8MPgA" } }, { "cell_type": "code", "source": [ "import io\n", "import os\n", "import re\n", "import gc\n", "import json\n", "import torch\n", "import pysrt\n", "import pickle\n", "import pathlib\n", "import warnings\n", "import whisperx\n", "import subprocess\n", "\n", "import numpy as np\n", "import pandas as pd\n", "\n", "from copy import deepcopy\n", "from pprint import pprint\n", "from tqdm.notebook import tqdm\n", "from collections import defaultdict, Counter\n", "from rapidfuzz import process, fuzz\n", "from sklearn.feature_extraction.text import TfidfVectorizer" ], "metadata": { "id": "wv5TA_UhXh1m" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "# Suppress pyannote warnings\n", "warnings.filterwarnings(\"ignore\", module=\"pyannote\")\n", "\n", "def build_set2_tfidf(set2, preferred_speaker_names):\n", " \"\"\"\n", " Projects TVQA speaker names to Whisper transcription and standardize them,\n", " using preferrewd TMDB names (first + last) -- they are preferred bc\n", " they represent (more) canonical forms of the names and are usually associated\n", " with external IDs like IMDb -- presents opportunities for further\n", " downstream analysis.\n", "\n", " set2:\n", " TVQA subtitles\n", "\n", " preferred_speaker_names:\n", " TMDb cast list names\n", " \"\"\"\n", " set2_texts = [seg.get('sentence_text', '') for seg in set2]\n", " vectorizer = TfidfVectorizer(norm='l2')\n", " set2_matrix = vectorizer.fit_transform(set2_texts)\n", " real_speakers_in_set2 = {\n", " seg['speaker']\n", " for seg in set2\n", " if seg.get('speaker') and seg['speaker'].lower() != 'unknown'\n", " }\n", " speaker_map = {}\n", " preferred_tokenized = {name: set(name.lower().split()) for name in preferred_speaker_names}\n", " for sp in real_speakers_in_set2:\n", " sp_tokens = set(sp.lower().split())\n", " candidates = [name for name, tokens in preferred_tokenized.items() if sp_tokens & tokens]\n", " if len(candidates) == 1:\n", " speaker_map[sp] = candidates[0]\n", " elif candidates:\n", " best = process.extractOne(sp, candidates, scorer=fuzz.ratio)\n", " if best:\n", " speaker_map[sp] = best[0]\n", " else:\n", " speaker_map[sp] = sp\n", " return vectorizer, set2_matrix, speaker_map\n", "\n", "def align_set1_with_tfidf(set1, set2, vectorizer, set2_matrix, speaker_map):\n", " \"\"\"\n", " Aligns TVQA subtitles with Whisper transcriptions, based on algorithm 1 in:\n", "\n", " Qingqiu Huang, Yu Xiong, Anyi Rao, Jiaze Wang, and Dahua Lin. \"MovieNet: A Holistic Dataset for Movie understanding\" [ECCV 2020]\n", "\n", " Args:\n", " set1: the texts to be aligned to, like whisper segments\n", " set2: the texts with more info, like original TVQA subtitles\n", " Returns: a new list of dict for set1 with improved 'speaker'.\n", " \"\"\"\n", " index_to_speaker = [seg.get('speaker', 'unknown') for seg in set2]\n", "\n", " set1_texts = [seg.get('text', '') for seg in set1]\n", " set1_matrix = vectorizer.transform(set1_texts) # shape (N, vocab_size)\n", " sim_matrix = set1_matrix.dot(set2_matrix.T) # shape (N, M)\n", " sim_array = sim_matrix.toarray()\n", " best_indices = np.argmax(sim_array, axis=1) # shape (N,)\n", "\n", " # First pass:\n", " index_to_speaker = [seg.get('speaker', 'unknown') for seg in set2]\n", " set1_texts = [seg.get('text', '') for seg in set1]\n", " set1_matrix = vectorizer.transform(set1_texts)\n", " sim_matrix = set1_matrix.dot(set2_matrix.T)\n", " sim_array = sim_matrix.toarray()\n", " best_indices = np.argmax(sim_array, axis=1)\n", "\n", " updated_set1 = []\n", " for i, seg1 in enumerate(set1):\n", " j = best_indices[i]\n", " candidate_spk = index_to_speaker[j] if 0 <= j < len(index_to_speaker) else 'unknown'\n", " seg_copy = dict(seg1)\n", " seg_copy['speaker'] = candidate_spk\n", " updated_set1.append(seg_copy)\n", "\n", " # Second pass:\n", "\n", " # The general idea is that we want to increase recall and get as many names\n", " # as possible, because they will all be manually fixed later.\n", " # So, we want to reduce the number of `unknown` speakers.\n", " # For each segment that ended withunknown, check:\n", " # if we previously found a real speaker for that same generic label\n", " # in an earlier segment:\n", " # if so, reuse it :)\n", "\n", " second_pass_result = []\n", " last_known_for_label = {}\n", " for seg_orig, seg_upd in zip(set1, updated_set1):\n", " gen_label = seg_orig.get('speaker')\n", " assigned_spk = seg_upd.get('speaker', 'unknown') or 'unknown'\n", " if assigned_spk.lower() != 'unknown':\n", " last_known_for_label[gen_label] = assigned_spk\n", " second_pass_result.append(seg_upd)\n", " else:\n", " if gen_label in last_known_for_label:\n", " seg_upd['speaker'] = last_known_for_label[gen_label]\n", " second_pass_result.append(seg_upd)\n", "\n", " for seg in second_pass_result:\n", " spk = seg.get('speaker', 'unknown')\n", " if spk.lower() != 'unknown' and spk in speaker_map:\n", " seg['speaker'] = speaker_map[spk]\n", "\n", " return second_pass_result\n", "\n", "def get_whisperx_segments(audio_file, model, diarize_model, device=\"cpu\", batch_size=16):\n", " audio = whisperx.load_audio(audio_file)\n", " result = model.transcribe(audio, batch_size=batch_size)\n", " model_a, metadata = whisperx.load_align_model(language_code=result[\"language\"], device=device)\n", " result = whisperx.align(result[\"segments\"], model_a, metadata, audio, device, return_char_alignments=False)\n", " diarize_segments = diarize_model(audio)\n", " result = whisperx.assign_word_speakers(diarize_segments, result)\n", " gc.collect()\n", " torch.cuda.empty_cache()\n", " return result['segments']\n", "\n", "def load_json(path):\n", " with open(path) as o:\n", " return json.load(o)\n", "\n", "def print_set(dataset):\n", " for seg in dataset:\n", " pprint({k: v for k, v in seg.items() if k != 'words'})" ], "metadata": { "id": "sz6V1cEtWf67" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "# Ensure you are using a GPU runtime in Colab (Runtime -> Change runtime type)\n", "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", "\n", "# Use \"int8\" if CPU-only\n", "compute_type = \"float16\"\n", "\n", "print(f\"Loading models to device: {device}\")\n", "\n", "# Load Whisper model\n", "model = whisperx.load_model(\n", " \"large-v2\",\n", " device=device,\n", " compute_type=compute_type,\n", " language='en'\n", ")\n", "\n", "# Load Diarization model\n", "diarize_model = whisperx.DiarizationPipeline(use_auth_token=True, device=device)" ], "metadata": { "id": "szy7MrmcWilH" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "source": [ "---" ], "metadata": { "id": "CKFZngkiZNOK" } }, { "cell_type": "markdown", "source": [ "## Transcribe with WhisperX" ], "metadata": { "id": "mQqQ81D9bObE" } }, { "cell_type": "code", "source": [ "mp3_file = 's01e02_seg02_clip_01.mp3'\n", "srt_file = 's01e02_seg02_clip_01.srt'\n", "annotation_file = 's01e02_seg02_clip_01.annotation.json'" ], "metadata": { "id": "Pz8MaHNLZLf3" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "# Step 1: Transcribe Audio with WhisperX to create `set1`\n", "\n", "# Good to run first to make sure your WhisperX works\n", "set1 = get_whisperx_segments(mp3_file, model, diarize_model, device)" ], "metadata": { "id": "fRLT1nlJXuq5" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "# Step 2: Prepare ground-truth `set2` and speaker names\n", "\n", "# Load preferred speaker names from the annotation file\n", "annotations = load_json(annotation_file)\n", "preferred_speaker_names = list(set(item['speaker'] for item in annotations))\n", "\n", "# Parse the SRT file to create `set2`\n", "set2 = []\n", "subs = pysrt.open(srt_file)\n", "for sub in subs:\n", " text = sub.text.replace('\\n', ' ').replace('\\r', ' ').strip()\n", " speaker_match = re.search(r'\\((.*?)\\:?\\)', text)\n", " speaker = speaker_match.group(1) if speaker_match else \"unknown\"\n", " cleaned_text = re.sub(r'\\(.*?\\)', '', text).replace('...', '').strip()\n", "\n", " # This block handles splitting lines as seen in the original script\n", " if ' - ' in cleaned_text:\n", " parts = cleaned_text.lstrip(\"- \").split(\" - \")\n", " for part in parts:\n", " part = part.strip()\n", " if part:\n", " new_seg = {'speaker': speaker, 'sentence_text': part}\n", " set2.append(new_seg)\n", " else:\n", " set2.append({'speaker': speaker, 'sentence_text': cleaned_text})" ], "metadata": { "id": "aiM8Ma_GWlUM" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "# Step 3: Align `set1` and `set2`\n", "\n", "vectorizer, set2_matrix, speaker_map = build_set2_tfidf(set2, preferred_speaker_names)\n", "\n", "improved_set1 = align_set1_with_tfidf(\n", " set1,\n", " set2,\n", " vectorizer,\n", " set2_matrix,\n", " speaker_map\n", ")\n", "\n", "# Optional -- check results:\n", "# print(f\"Sanity check: final enhanced subtitle data (improved_set1)\")\n", "# print_set(improved_set1)" ], "metadata": { "id": "PaIo7jrCWoEk" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "source": [ "---" ], "metadata": { "id": "q3zJ-ub-aWi3" } }, { "cell_type": "markdown", "source": [ "## Add text to `annotation.json`" ], "metadata": { "id": "u4Z1kV8lbo7J" } }, { "cell_type": "code", "source": [ "final_output_file = 's01e02_seg02_clip_01.aligned_annotation.json'" ], "metadata": { "id": "P61j_LAqba0f" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "assert len(annotations) == len(improved_set1)\n", "\n", "final_data = []\n", "# Use zip() to pair each original annotation with its corresponding transcribed segment\n", "for annotation_item, transcribed_item in zip(annotations, improved_set1):\n", " # Copy the original item to preserve all its metadata (speaker, etc.)\n", " new_item = annotation_item.copy()\n", "\n", " # Add the 'text' from the transcribed item\n", " new_item['text'] = transcribed_item.get('text', '').strip()\n", "\n", " final_data.append(new_item)\n", "\n", "with open(final_output_file, 'w') as o:\n", " json.dump(final_data, o, indent=2)\n", "\n", "print(f\"Aligned file saved to: '{final_output_file}'\")\n", "pprint(final_data[:3])" ], "metadata": { "id": "iepxoujTZ7m7" }, "execution_count": null, "outputs": [] } ] }