VwArship.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "j9TI-Q6m3qlx",
        "cellView": "form"
      },
      "source": [
        "#@markdown **Mount Google Drive** (skip this if your audio file isn't stored there)\n",
        "from google.colab import drive\n",
        "drive.mount('drive')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "teF-Ut8Z7Gjp",
        "cellView": "form"
      },
      "source": [
        "#@markdown **Setup Whisper**\n",
        "!apt-get install sox libsox-fmt-mp3 libsndfile1 ffmpeg\n",
        "!pip install deepl srt ffmpeg-python spleeter\n",
        "!pip install git+https://github.com/openai/whisper.git\n",
        "print(\"Done\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sos9vsxPkIN7",
        "cellView": "form"
      },
      "source": [
        "#@markdown **Run Whisper**\n",
        "\n",
        "# @markdown Required settings:\n",
        "audio_path = \"input.mp3\"  # @param {type:\"string\"}\n",
        "model_size = \"large\"  # @param [\"tiny\",\"base\",\"small\",\"medium\", \"large\"]\n",
        "language = \"japanese\"  # @param {type:\"string\"}\n",
        "translation_mode = \"End-to-end Whisper (default)\"  # @param [\"End-to-end Whisper (default)\", \"Whisper -> DeepL\", \"No translation\"]\n",
        "# @markdown Advanced settings:\n",
        "deepl_authkey = \"\"  # @param {type:\"string\"}\n",
        "source_separation = False  # @param {type:\"boolean\"}\n",
        "vad_threshold = 0.4  # @param {type:\"number\"}\n",
        "chunk_threshold = 3.0  # @param {type:\"number\"}\n",
        "deepl_target_lang = \"EN-US\"  # @param {type:\"string\"}\n",
        "max_attempts = 1  # @param {type:\"integer\"}\n",
        "initial_prompt = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "import tensorflow as tf\n",
        "import torch\n",
        "import whisper\n",
        "import os\n",
        "import ffmpeg\n",
        "import srt\n",
        "from tqdm import tqdm\n",
        "import datetime\n",
        "import deepl\n",
        "import urllib.request\n",
        "import json\n",
        "from google.colab import files\n",
        "\n",
        "# Configuration\n",
        "assert max_attempts >= 1\n",
        "assert vad_threshold >= 0.01\n",
        "assert chunk_threshold >= 0.1\n",
        "assert audio_path != \"\"\n",
        "assert language != \"\"\n",
        "if translation_mode == \"End-to-end Whisper (default)\":\n",
        "    task = \"translate\"\n",
        "    run_deepl = False\n",
        "elif translation_mode == \"Whisper -> DeepL\":\n",
        "    task = \"transcribe\"\n",
        "    run_deepl = True\n",
        "elif translation_mode == \"No translation\":\n",
        "    task = \"transcribe\"\n",
        "    run_deepl = False\n",
        "else:\n",
        "    raise ValueError(\"Invalid translation mode\")\n",
        "if initial_prompt.strip() == \"\":\n",
        "    initial_prompt = None\n",
        "\n",
        "if \"http://\" in audio_path or \"https://\" in audio_path:\n",
        "    print(\"Downloading audio...\")\n",
        "    urllib.request.urlretrieve(audio_path, \"input_file\")\n",
        "    audio_path = \"input_file\"\n",
        "else:\n",
        "    if not os.path.exists(audio_path):\n",
        "        try:\n",
        "            audio_path = uploaded_file\n",
        "            if not os.path.exists(audio_path):\n",
        "                raise ValueError(\"Input audio not found. Is your audio_path correct?\")\n",
        "        except NameError:\n",
        "            raise ValueError(\"Input audio not found. Did you upload a file?\")\n",
        "\n",
        "out_path = os.path.splitext(audio_path)[0] + \".srt\"\n",
        "out_path_pre = os.path.splitext(audio_path)[0] + \"_Untranslated.srt\"\n",
        "if source_separation:\n",
        "    print(\"Separating vocals...\")\n",
        "    !ffprobe -i \"{audio_path}\" -show_entries format=duration -v quiet -of csv=\"p=0\" > input_length\n",
        "    with open(\"input_length\") as f:\n",
        "        input_length = int(float(f.read())) + 1\n",
        "    !spleeter separate -d {input_length} -p spleeter:2stems -o output \"{audio_path}\"\n",
        "    spleeter_dir = os.path.basename(os.path.splitext(audio_path)[0])\n",
        "    audio_path = \"output/\" + spleeter_dir + \"/vocals.wav\"\n",
        "\n",
        "print(\"Encoding audio...\")\n",
        "if not os.path.exists(\"vad_chunks\"):\n",
        "    os.mkdir(\"vad_chunks\")\n",
        "ffmpeg.input(audio_path).output(\n",
        "    \"vad_chunks/silero_temp.wav\",\n",
        "    ar=\"16000\",\n",
        "    ac=\"1\",\n",
        "    acodec=\"pcm_s16le\",\n",
        "    map_metadata=\"-1\",\n",
        "    fflags=\"+bitexact\",\n",
        ").overwrite_output().run(quiet=True)\n",
        "\n",
        "print(\"Running VAD...\")\n",
        "model, utils = torch.hub.load(\n",
        "    repo_or_dir=\"snakers4/silero-vad\", model=\"silero_vad\", onnx=False\n",
        ")\n",
        "\n",
        "(get_speech_timestamps, save_audio, read_audio, VADIterator, collect_chunks) = utils\n",
        "\n",
        "# Generate VAD timestamps\n",
        "VAD_SR = 16000\n",
        "wav = read_audio(\"vad_chunks/silero_temp.wav\", sampling_rate=VAD_SR)\n",
        "t = get_speech_timestamps(wav, model, sampling_rate=VAD_SR, threshold=vad_threshold)\n",
        "\n",
        "# Add a bit of padding, and remove small gaps\n",
        "for i in range(len(t)):\n",
        "    t[i][\"start\"] = max(0, t[i][\"start\"] - 3200)  # 0.2s head\n",
        "    t[i][\"end\"] = min(wav.shape[0] - 16, t[i][\"end\"] + 20800)  # 1.3s tail\n",
        "    if i > 0 and t[i][\"start\"] < t[i - 1][\"end\"]:\n",
        "        t[i][\"start\"] = t[i - 1][\"end\"]  # Remove overlap\n",
        "\n",
        "# If breaks are longer than chunk_threshold seconds, split into a new audio file\n",
        "# This'll effectively turn long transcriptions into many shorter ones\n",
        "u = [[]]\n",
        "for i in range(len(t)):\n",
        "    if i > 0 and t[i][\"start\"] > t[i - 1][\"end\"] + (chunk_threshold * VAD_SR):\n",
        "        u.append([])\n",
        "    u[-1].append(t[i])\n",
        "\n",
        "# Merge speech chunks\n",
        "for i in range(len(u)):\n",
        "    save_audio(\n",
        "        \"vad_chunks/\" + str(i) + \".wav\",\n",
        "        collect_chunks(u[i], wav),\n",
        "        sampling_rate=VAD_SR,\n",
        "    )\n",
        "os.remove(\"vad_chunks/silero_temp.wav\")\n",
        "\n",
        "# Convert timestamps to seconds\n",
        "for i in range(len(u)):\n",
        "    time = 0.0\n",
        "    offset = 0.0\n",
        "    for j in range(len(u[i])):\n",
        "        u[i][j][\"start\"] /= VAD_SR\n",
        "        u[i][j][\"end\"] /= VAD_SR\n",
        "        u[i][j][\"chunk_start\"] = time\n",
        "        time += u[i][j][\"end\"] - u[i][j][\"start\"]\n",
        "        u[i][j][\"chunk_end\"] = time\n",
        "        if j == 0:\n",
        "            offset += u[i][j][\"start\"]\n",
        "        else:\n",
        "            offset += u[i][j][\"start\"] - u[i][j - 1][\"end\"]\n",
        "        u[i][j][\"offset\"] = offset\n",
        "\n",
        "# Run Whisper on each audio chunk\n",
        "print(\"Running Whisper...\")\n",
        "model = whisper.load_model(model_size)\n",
        "subs = []\n",
        "segment_info = []\n",
        "sub_index = 1\n",
        "suppress_low = [\n",
        "    \"Thank you\",\n",
        "    \"Thanks for\",\n",
        "    \"ike and \",\n",
        "    \"Bye.\",\n",
        "    \"Bye!\",\n",
        "    \"Bye bye!\",\n",
        "    \"lease sub\",\n",
        "    \"The end.\",\n",
        "    \"視聴\",\n",
        "]\n",
        "suppress_high = [\n",
        "    \"ubscribe\",\n",
        "    \"my channel\",\n",
        "    \"the channel\",\n",
        "    \"our channel\",\n",
        "    \"ollow me on\",\n",
        "    \"for watching\",\n",
        "    \"hank you for watching\",\n",
        "    \"for your viewing\",\n",
        "    \"r viewing\",\n",
        "    \"Amara\",\n",
        "    \"next video\",\n",
        "    \"full video\",\n",
        "    \"ranslation by\",\n",
        "    \"ranslated by\",\n",
        "    \"ee you next week\",\n",
        "    \"ご視聴\",\n",
        "    \"視聴ありがとうございました\",\n",
        "]\n",
        "for i in tqdm(range(len(u))):\n",
        "    line_buffer = []  # Used for DeepL\n",
        "    for x in range(max_attempts):\n",
        "        result = model.transcribe(\n",
        "            \"vad_chunks/\" + str(i) + \".wav\", task=task, language=language, initial_prompt=initial_prompt\n",
        "        )\n",
        "        # Break if result doesn't end with severe hallucinations\n",
        "        if len(result[\"segments\"]) == 0:\n",
        "            break\n",
        "        elif result[\"segments\"][-1][\"end\"] < u[i][-1][\"chunk_end\"] + 10.0:\n",
        "            break\n",
        "        elif x+1 < max_attempts:\n",
        "            print(\"Retrying chunk\", i)\n",
        "    for r in result[\"segments\"]:\n",
        "        # Skip audio timestamped after the chunk has ended\n",
        "        if r[\"start\"] > u[i][-1][\"chunk_end\"]:\n",
        "            continue\n",
        "        # Reduce log probability for certain words/phrases\n",
        "        for s in suppress_low:\n",
        "            if s in r[\"text\"]:\n",
        "                r[\"avg_logprob\"] -= 0.15\n",
        "        for s in suppress_high:\n",
        "            if s in r[\"text\"]:\n",
        "                r[\"avg_logprob\"] -= 0.35\n",
        "        # Keep segment info for debugging\n",
        "        del r[\"tokens\"]\n",
        "        segment_info.append(r)\n",
        "        # Skip if log prob is low or no speech prob is high\n",
        "        if r[\"avg_logprob\"] < -1.0 or r[\"no_speech_prob\"] > 0.7:\n",
        "            continue\n",
        "        # Set start timestamp\n",
        "        start = r[\"start\"] + u[i][0][\"offset\"]\n",
        "        for j in range(len(u[i])):\n",
        "            if (\n",
        "                r[\"start\"] >= u[i][j][\"chunk_start\"]\n",
        "                and r[\"start\"] <= u[i][j][\"chunk_end\"]\n",
        "            ):\n",
        "                start = r[\"start\"] + u[i][j][\"offset\"]\n",
        "                break\n",
        "        # Prevent overlapping subs\n",
        "        if len(subs) > 0:\n",
        "            last_end = datetime.timedelta.total_seconds(subs[-1].end)\n",
        "            if last_end > start:\n",
        "                subs[-1].end = datetime.timedelta(seconds=start)\n",
        "        # Set end timestamp\n",
        "        end = u[i][-1][\"end\"] + 0.5\n",
        "        for j in range(len(u[i])):\n",
        "            if r[\"end\"] >= u[i][j][\"chunk_start\"] and r[\"end\"] <= u[i][j][\"chunk_end\"]:\n",
        "                end = r[\"end\"] + u[i][j][\"offset\"]\n",
        "                break\n",
        "        # Add to SRT list\n",
        "        subs.append(\n",
        "            srt.Subtitle(\n",
        "                index=sub_index,\n",
        "                start=datetime.timedelta(seconds=start),\n",
        "                end=datetime.timedelta(seconds=end),\n",
        "                content=r[\"text\"].strip(),\n",
        "            )\n",
        "        )\n",
        "        sub_index += 1\n",
        "\n",
        "with open(\"segment_info.json\", \"w\", encoding=\"utf8\") as f:\n",
        "    json.dump(segment_info, f, indent=4)\n",
        "\n",
        "# DeepL translation\n",
        "translate_error = False\n",
        "if run_deepl:\n",
        "    print(\"Translating...\")\n",
        "    with open(out_path_pre, \"w\", encoding=\"utf8\") as f:\n",
        "        f.write(srt.compose(subs))\n",
        "    print(\"(Untranslated subs saved to\", out_path_pre, \")\")\n",
        "\n",
        "    lines = []\n",
        "    punct_match = [\"。\", \"、\", \",\", \".\", \"〜\", \"！\", \"!\", \"？\", \"?\", \"-\"]\n",
        "    for i in range(len(subs)):\n",
        "        if language.lower() == \"japanese\":\n",
        "            if subs[i].content[-1] not in punct_match:\n",
        "                subs[i].content += \"。\"\n",
        "            subs[i].content = \"「\" + subs[i].content + \"」\"\n",
        "        else:\n",
        "            if subs[i].content[-1] not in punct_match:\n",
        "                subs[i].content += \".\"\n",
        "            subs[i].content = '\"' + subs[i].content + '\"'\n",
        "    for i in range(len(subs)):\n",
        "        lines.append(subs[i].content)\n",
        "\n",
        "    grouped_lines = []\n",
        "    english_lines = []\n",
        "    for i, l in enumerate(lines):\n",
        "        if i % 30 == 0:\n",
        "            # Split lines into smaller groups, to prevent error 413\n",
        "            grouped_lines.append([])\n",
        "            if i != 0:\n",
        "                # Include previous 3 lines, to preserve context between splits\n",
        "                grouped_lines[-1].extend(grouped_lines[-2][-3:])\n",
        "        grouped_lines[-1].append(l.strip())\n",
        "        \n",
        "    try:\n",
        "        translator = deepl.Translator(deepl_authkey)\n",
        "        for i, n in enumerate(tqdm(grouped_lines)):\n",
        "            x = [\"\\n\".join(n).strip()]\n",
        "            if language.lower() == \"japanese\":\n",
        "                result = translator.translate_text(x, source_lang=\"JA\", target_lang=deepl_target_lang)\n",
        "            else:\n",
        "                result = translator.translate_text(x, target_lang=deepl_target_lang)\n",
        "            english_tl = result[0].text.strip().splitlines()\n",
        "            assert len(english_tl) == len(n), (\n",
        "                \"Invalid translation line count (\"\n",
        "                + str(len(english_tl))\n",
        "                + \" vs \"\n",
        "                + str(len(n))\n",
        "                + \")\"\n",
        "            )\n",
        "            if i != 0:\n",
        "                english_tl = english_tl[3:]\n",
        "            remove_quotes = dict.fromkeys(map(ord, '\"„“‟”＂「」'), None)\n",
        "            for e in english_tl:\n",
        "                english_lines.append(\n",
        "                    e.strip().translate(remove_quotes).replace(\"’\", \"'\")\n",
        "                )\n",
        "        for i, e in enumerate(english_lines):\n",
        "            subs[i].content = e\n",
        "    except Exception as e:\n",
        "        print(\"DeepL translation error:\", e)\n",
        "        print(\"(downloading untranslated version instead)\")\n",
        "        translate_error = True\n",
        "\n",
        "# Write SRT file\n",
        "if translate_error:\n",
        "    files.download(out_path_pre)\n",
        "else:\n",
        "    # Removal of garbage lines\n",
        "    garbage_list = [\n",
        "        \"a\",\n",
        "        \"aa\",\n",
        "        \"ah\",\n",
        "        \"ahh\",\n",
        "        \"ha\",\n",
        "        \"haa\",\n",
        "        \"hah\",\n",
        "        \"haha\",\n",
        "        \"hahaha\",\n",
        "        \"mmm\",\n",
        "        \"mm\",\n",
        "        \"m\",\n",
        "        \"h\",\n",
        "        \"o\",\n",
        "        \"mh\",\n",
        "        \"mmh\",\n",
        "        \"hm\",\n",
        "        \"hmm\",\n",
        "        \"huh\",\n",
        "        \"oh\",\n",
        "    ]\n",
        "    need_context_lines = [\n",
        "        \"feelsgod\",\n",
        "        \"godbye\",\n",
        "        \"godnight\",\n",
        "        \"thankyou\",\n",
        "    ]\n",
        "    clean_subs = list()\n",
        "    last_line_garbage = False\n",
        "    for i in range(len(subs)):\n",
        "        c = subs[i].content\n",
        "        c = (\n",
        "            c.replace(\".\", \"\")\n",
        "            .replace(\",\", \"\")\n",
        "            .replace(\":\", \"\")\n",
        "            .replace(\";\", \"\")\n",
        "            .replace(\"!\", \"\")\n",
        "            .replace(\"?\", \"\")\n",
        "            .replace(\"-\", \" \")\n",
        "            .replace(\"  \", \" \")\n",
        "            .replace(\"  \", \" \")\n",
        "            .replace(\"  \", \" \")\n",
        "            .lower()\n",
        "            .replace(\"that feels\", \"feels\")\n",
        "            .replace(\"it feels\", \"feels\")\n",
        "            .replace(\"feels good\", \"feelsgood\")\n",
        "            .replace(\"good bye\", \"goodbye\")\n",
        "            .replace(\"good night\", \"goodnight\")\n",
        "            .replace(\"thank you\", \"thankyou\")\n",
        "            .replace(\"aaaaaa\", \"a\")\n",
        "            .replace(\"aaaa\", \"a\")\n",
        "            .replace(\"aa\", \"a\")\n",
        "            .replace(\"aa\", \"a\")\n",
        "            .replace(\"mmmmmm\", \"m\")\n",
        "            .replace(\"mmmm\", \"m\")\n",
        "            .replace(\"mm\", \"m\")\n",
        "            .replace(\"mm\", \"m\")\n",
        "            .replace(\"hhhhhh\", \"h\")\n",
        "            .replace(\"hhhh\", \"h\")\n",
        "            .replace(\"hh\", \"h\")\n",
        "            .replace(\"hh\", \"h\")\n",
        "            .replace(\"oooooo\", \"o\")\n",
        "            .replace(\"oooo\", \"o\")\n",
        "            .replace(\"oo\", \"o\")\n",
        "            .replace(\"oo\", \"o\")\n",
        "        )\n",
        "        is_garbage = True\n",
        "        for w in c.split(\" \"):\n",
        "            if w.strip() == \"\":\n",
        "                continue\n",
        "            if w.strip() in garbage_list:\n",
        "                continue\n",
        "            elif w.strip() in need_context_lines and last_line_garbage:\n",
        "                continue\n",
        "            else:\n",
        "                is_garbage = False\n",
        "                break\n",
        "        if not is_garbage:\n",
        "            clean_subs.append(subs[i])\n",
        "        last_line_garbage = is_garbage\n",
        "    with open(out_path, \"w\", encoding=\"utf8\") as f:\n",
        "        f.write(srt.compose(clean_subs))\n",
        "    print(\"\\nDone! Subs written to\", out_path)\n",
        "    print(\"Downloading SRT file:\")\n",
        "    files.download(out_path)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
        {
      "cell_type": "code",
      "metadata": {
        "id": "sos9vsxPkIN7",
        "cellView": "form"
      },
      "source": [
        "#@markdown **Run Whisper**\n",
        "\n",
        "# @markdown Required settings:\n",
        "audio_path = \"input.mp3\"  # @param {type:\"string\"}\n",
        "model_size = \"large\"  # @param [\"tiny\",\"base\",\"small\",\"medium\", \"large\"]\n",
        "language = \"japanese\"  # @param {type:\"string\"}\n",
        "translation_mode = \"End-to-end Whisper (default)\"  # @param [\"End-to-end Whisper (default)\", \"Whisper -> DeepL\", \"No translation\"]\n",
        "# @markdown Advanced settings:\n",
        "deepl_authkey = \"\"  # @param {type:\"string\"}\n",
        "source_separation = False  # @param {type:\"boolean\"}\n",
        "vad_threshold = 0.4  # @param {type:\"number\"}\n",
        "chunk_threshold = 3.0  # @param {type:\"number\"}\n",
        "deepl_target_lang = \"EN-US\"  # @param {type:\"string\"}\n",
        "max_attempts = 1  # @param {type:\"integer\"}\n",
        "initial_prompt = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "import tensorflow as tf\n",
        "import torch\n",
        "import whisper\n",
        "import os\n",
        "import ffmpeg\n",
        "import srt\n",
        "from tqdm import tqdm\n",
        "import datetime\n",
        "import deepl\n",
        "import urllib.request\n",
        "import json\n",
        "from google.colab import files\n",
        "\n",
        "# Configuration\n",
        "assert max_attempts >= 1\n",
        "assert vad_threshold >= 0.01\n",
        "assert chunk_threshold >= 0.1\n",
        "assert audio_path != \"\"\n",
        "assert language != \"\"\n",
        "if translation_mode == \"End-to-end Whisper (default)\":\n",
        "    task = \"translate\"\n",
        "    run_deepl = False\n",
        "elif translation_mode == \"Whisper -> DeepL\":\n",
        "    task = \"transcribe\"\n",
        "    run_deepl = True\n",
        "elif translation_mode == \"No translation\":\n",
        "    task = \"transcribe\"\n",
        "    run_deepl = False\n",
        "else:\n",
        "    raise ValueError(\"Invalid translation mode\")\n",
        "if initial_prompt.strip() == \"\":\n",
        "    initial_prompt = None\n",
        "\n",
        "if \"http://\" in audio_path or \"https://\" in audio_path:\n",
        "    print(\"Downloading audio...\")\n",
        "    urllib.request.urlretrieve(audio_path, \"input_file\")\n",
        "    audio_path = \"input_file\"\n",
        "else:\n",
        "    if not os.path.exists(audio_path):\n",
        "        try:\n",
        "            audio_path = uploaded_file\n",
        "            if not os.path.exists(audio_path):\n",
        "                raise ValueError(\"Input audio not found. Is your audio_path correct?\")\n",
        "        except NameError:\n",
        "            raise ValueError(\"Input audio not found. Did you upload a file?\")\n",
        "\n",
        "out_path = os.path.splitext(audio_path)[0] + \".srt\"\n",
        "out_path_pre = os.path.splitext(audio_path)[0] + \"_Untranslated.srt\"\n",
        "if source_separation:\n",
        "    print(\"Separating vocals...\")\n",
        "    !ffprobe -i \"{audio_path}\" -show_entries format=duration -v quiet -of csv=\"p=0\" > input_length\n",
        "    with open(\"input_length\") as f:\n",
        "        input_length = int(float(f.read())) + 1\n",
        "    !spleeter separate -d {input_length} -p spleeter:2stems -o output \"{audio_path}\"\n",
        "    spleeter_dir = os.path.basename(os.path.splitext(audio_path)[0])\n",
        "    audio_path = \"output/\" + spleeter_dir + \"/vocals.wav\"\n",
        "\n",
        "print(\"Encoding audio...\")\n",
        "if not os.path.exists(\"vad_chunks\"):\n",
        "    os.mkdir(\"vad_chunks\")\n",
        "ffmpeg.input(audio_path).output(\n",
        "    \"vad_chunks/silero_temp.wav\",\n",
        "    ar=\"16000\",\n",
        "    ac=\"1\",\n",
        "    acodec=\"pcm_s16le\",\n",
        "    map_metadata=\"-1\",\n",
        "    fflags=\"+bitexact\",\n",
        ").overwrite_output().run(quiet=True)\n",
        "\n",
        "print(\"Running VAD...\")\n",
        "model, utils = torch.hub.load(\n",
        "    repo_or_dir=\"snakers4/silero-vad\", model=\"silero_vad\", onnx=False\n",
        ")\n",
        "\n",
        "(get_speech_timestamps, save_audio, read_audio, VADIterator, collect_chunks) = utils\n",
        "\n",
        "# Generate VAD timestamps\n",
        "VAD_SR = 16000\n",
        "wav = read_audio(\"vad_chunks/silero_temp.wav\", sampling_rate=VAD_SR)\n",
        "t = get_speech_timestamps(wav, model, sampling_rate=VAD_SR, threshold=vad_threshold)\n",
        "\n",
        "# Add a bit of padding, and remove small gaps\n",
        "for i in range(len(t)):\n",
        "    t[i][\"start\"] = max(0, t[i][\"start\"] - 3200)  # 0.2s head\n",
        "    t[i][\"end\"] = min(wav.shape[0] - 16, t[i][\"end\"] + 20800)  # 1.3s tail\n",
        "    if i > 0 and t[i][\"start\"] < t[i - 1][\"end\"]:\n",
        "        t[i][\"start\"] = t[i - 1][\"end\"]  # Remove overlap\n",
        "\n",
        "# If breaks are longer than chunk_threshold seconds, split into a new audio file\n",
        "# This'll effectively turn long transcriptions into many shorter ones\n",
        "u = [[]]\n",
        "for i in range(len(t)):\n",
        "    if i > 0 and t[i][\"start\"] > t[i - 1][\"end\"] + (chunk_threshold * VAD_SR):\n",
        "        u.append([])\n",
        "    u[-1].append(t[i])\n",
        "\n",
        "# Merge speech chunks\n",
        "for i in range(len(u)):\n",
        "    save_audio(\n",
        "        \"vad_chunks/\" + str(i) + \".wav\",\n",
        "        collect_chunks(u[i], wav),\n",
        "        sampling_rate=VAD_SR,\n",
        "    )\n",
        "os.remove(\"vad_chunks/silero_temp.wav\")\n",
        "\n",
        "# Convert timestamps to seconds\n",
        "for i in range(len(u)):\n",
        "    time = 0.0\n",
        "    offset = 0.0\n",
        "    for j in range(len(u[i])):\n",
        "        u[i][j][\"start\"] /= VAD_SR\n",
        "        u[i][j][\"end\"] /= VAD_SR\n",
        "        u[i][j][\"chunk_start\"] = time\n",
        "        time += u[i][j][\"end\"] - u[i][j][\"start\"]\n",
        "        u[i][j][\"chunk_end\"] = time\n",
        "        if j == 0:\n",
        "            offset += u[i][j][\"start\"]\n",
        "        else:\n",
        "            offset += u[i][j][\"start\"] - u[i][j - 1][\"end\"]\n",
        "        u[i][j][\"offset\"] = offset\n",
        "\n",
        "# Run Whisper on each audio chunk\n",
        "print(\"Running Whisper...\")\n",
        "model = whisper.load_model(model_size)\n",
        "subs = []\n",
        "segment_info = []\n",
        "sub_index = 1\n",
        "suppress_low = [\n",
        "    \"Thank you\",\n",
        "    \"Thanks for\",\n",
        "    \"ike and \",\n",
        "    \"Bye.\",\n",
        "    \"Bye!\",\n",
        "    \"Bye bye!\",\n",
        "    \"lease sub\",\n",
        "    \"The end.\",\n",
        "    \"視聴\",\n",
        "]\n",
        "suppress_high = [\n",
        "    \"ubscribe\",\n",
        "    \"my channel\",\n",
        "    \"the channel\",\n",
        "    \"our channel\",\n",
        "    \"ollow me on\",\n",
        "    \"for watching\",\n",
        "    \"hank you for watching\",\n",
        "    \"for your viewing\",\n",
        "    \"r viewing\",\n",
        "    \"Amara\",\n",
        "    \"next video\",\n",
        "    \"full video\",\n",
        "    \"ranslation by\",\n",
        "    \"ranslated by\",\n",
        "    \"ee you next week\",\n",
        "    \"ご視聴\",\n",
        "    \"視聴ありがとうございました\",\n",
        "]\n",
        "for i in tqdm(range(len(u))):\n",
        "    line_buffer = []  # Used for DeepL\n",
        "    for x in range(max_attempts):\n",
        "        result = model.transcribe(\n",
        "            \"vad_chunks/\" + str(i) + \".wav\", task=task, language=language, initial_prompt=initial_prompt\n",
        "        )\n",
        "        # Break if result doesn't end with severe hallucinations\n",
        "        if len(result[\"segments\"]) == 0:\n",
        "            break\n",
        "        elif result[\"segments\"][-1][\"end\"] < u[i][-1][\"chunk_end\"] + 10.0:\n",
        "            break\n",
        "        elif x+1 < max_attempts:\n",
        "            print(\"Retrying chunk\", i)\n",
        "    for r in result[\"segments\"]:\n",
        "        # Skip audio timestamped after the chunk has ended\n",
        "        if r[\"start\"] > u[i][-1][\"chunk_end\"]:\n",
        "            continue\n",
        "        # Reduce log probability for certain words/phrases\n",
        "        for s in suppress_low:\n",
        "            if s in r[\"text\"]:\n",
        "                r[\"avg_logprob\"] -= 0.15\n",
        "        for s in suppress_high:\n",
        "            if s in r[\"text\"]:\n",
        "                r[\"avg_logprob\"] -= 0.35\n",
        "        # Keep segment info for debugging\n",
        "        del r[\"tokens\"]\n",
        "        segment_info.append(r)\n",
        "        # Skip if log prob is low or no speech prob is high\n",
        "        if r[\"avg_logprob\"] < -1.0 or r[\"no_speech_prob\"] > 0.7:\n",
        "            continue\n",
        "        # Set start timestamp\n",
        "        start = r[\"start\"] + u[i][0][\"offset\"]\n",
        "        for j in range(len(u[i])):\n",
        "            if (\n",
        "                r[\"start\"] >= u[i][j][\"chunk_start\"]\n",
        "                and r[\"start\"] <= u[i][j][\"chunk_end\"]\n",
        "            ):\n",
        "                start = r[\"start\"] + u[i][j][\"offset\"]\n",
        "                break\n",
        "        # Prevent overlapping subs\n",
        "        if len(subs) > 0:\n",
        "            last_end = datetime.timedelta.total_seconds(subs[-1].end)\n",
        "            if last_end > start:\n",
        "                subs[-1].end = datetime.timedelta(seconds=start)\n",
        "        # Set end timestamp\n",
        "        end = u[i][-1][\"end\"] + 0.5\n",
        "        for j in range(len(u[i])):\n",
        "            if r[\"end\"] >= u[i][j][\"chunk_start\"] and r[\"end\"] <= u[i][j][\"chunk_end\"]:\n",
        "                end = r[\"end\"] + u[i][j][\"offset\"]\n",
        "                break\n",
        "        # Add to SRT list\n",
        "        subs.append(\n",
        "            srt.Subtitle(\n",
        "                index=sub_index,\n",
        "                start=datetime.timedelta(seconds=start),\n",
        "                end=datetime.timedelta(seconds=end),\n",
        "                content=r[\"text\"].strip(),\n",
        "            )\n",
        "        )\n",
        "        sub_index += 1\n",
        "\n",
        "with open(\"segment_info.json\", \"w\", encoding=\"utf8\") as f:\n",
        "    json.dump(segment_info, f, indent=4)\n",
        "\n",
        "# DeepL translation\n",
        "translate_error = False\n",
        "if run_deepl:\n",
        "    print(\"Translating...\")\n",
        "    with open(out_path_pre, \"w\", encoding=\"utf8\") as f:\n",
        "        f.write(srt.compose(subs))\n",
        "    print(\"(Untranslated subs saved to\", out_path_pre, \")\")\n",
        "\n",
        "    lines = []\n",
        "    punct_match = [\"。\", \"、\", \",\", \".\", \"〜\", \"！\", \"!\", \"？\", \"?\", \"-\"]\n",
        "    for i in range(len(subs)):\n",
        "        if language.lower() == \"japanese\":\n",
        "            if subs[i].content[-1] not in punct_match:\n",
        "                subs[i].content += \"。\"\n",
        "            subs[i].content = \"「\" + subs[i].content + \"」\"\n",
        "        else:\n",
        "            if subs[i].content[-1] not in punct_match:\n",
        "                subs[i].content += \".\"\n",
        "            subs[i].content = '\"' + subs[i].content + '\"'\n",
        "    for i in range(len(subs)):\n",
        "        lines.append(subs[i].content)\n",
        "\n",
        "    grouped_lines = []\n",
        "    english_lines = []\n",
        "    for i, l in enumerate(lines):\n",
        "        if i % 30 == 0:\n",
        "            # Split lines into smaller groups, to prevent error 413\n",
        "            grouped_lines.append([])\n",
        "            if i != 0:\n",
        "                # Include previous 3 lines, to preserve context between splits\n",
        "                grouped_lines[-1].extend(grouped_lines[-2][-3:])\n",
        "        grouped_lines[-1].append(l.strip())\n",
        "        \n",
        "    try:\n",
        "        translator = deepl.Translator(deepl_authkey)\n",
        "        for i, n in enumerate(tqdm(grouped_lines)):\n",
        "            x = [\"\\n\".join(n).strip()]\n",
        "            if language.lower() == \"japanese\":\n",
        "                result = translator.translate_text(x, source_lang=\"JA\", target_lang=deepl_target_lang)\n",
        "            else:\n",
        "                result = translator.translate_text(x, target_lang=deepl_target_lang)\n",
        "            english_tl = result[0].text.strip().splitlines()\n",
        "            assert len(english_tl) == len(n), (\n",
        "                \"Invalid translation line count (\"\n",
        "                + str(len(english_tl))\n",
        "                + \" vs \"\n",
        "                + str(len(n))\n",
        "                + \")\"\n",
        "            )\n",
        "            if i != 0:\n",
        "                english_tl = english_tl[3:]\n",
        "            remove_quotes = dict.fromkeys(map(ord, '\"„“‟”＂「」'), None)\n",
        "            for e in english_tl:\n",
        "                english_lines.append(\n",
        "                    e.strip().translate(remove_quotes).replace(\"’\", \"'\")\n",
        "                )\n",
        "        for i, e in enumerate(english_lines):\n",
        "            subs[i].content = e\n",
        "    except Exception as e:\n",
        "        print(\"DeepL translation error:\", e)\n",
        "        print(\"(downloading untranslated version instead)\")\n",
        "        translate_error = True\n",
        "\n",
        "# Write SRT file\n",
        "if translate_error:\n",
        "    files.download(out_path_pre)\n",
        "else:\n",
        "    # Removal of garbage lines\n",
        "    garbage_list = [\n",
        "        \"a\",\n",
        "        \"aa\",\n",
        "        \"ah\",\n",
        "        \"ahh\",\n",
        "        \"ha\",\n",
        "        \"haa\",\n",
        "        \"hah\",\n",
        "        \"haha\",\n",
        "        \"hahaha\",\n",
        "        \"mmm\",\n",
        "        \"mm\",\n",
        "        \"m\",\n",
        "        \"h\",\n",
        "        \"o\",\n",
        "        \"mh\",\n",
        "        \"mmh\",\n",
        "        \"hm\",\n",
        "        \"hmm\",\n",
        "        \"huh\",\n",
        "        \"oh\",\n",
        "    ]\n",
        "    need_context_lines = [\n",
        "        \"feelsgod\",\n",
        "        \"godbye\",\n",
        "        \"godnight\",\n",
        "        \"thankyou\",\n",
        "    ]\n",
        "    clean_subs = list()\n",
        "    last_line_garbage = False\n",
        "    for i in range(len(subs)):\n",
        "        c = subs[i].content\n",
        "        c = (\n",
        "            c.replace(\".\", \"\")\n",
        "            .replace(\",\", \"\")\n",
        "            .replace(\":\", \"\")\n",
        "            .replace(\";\", \"\")\n",
        "            .replace(\"!\", \"\")\n",
        "            .replace(\"?\", \"\")\n",
        "            .replace(\"-\", \" \")\n",
        "            .replace(\"  \", \" \")\n",
        "            .replace(\"  \", \" \")\n",
        "            .replace(\"  \", \" \")\n",
        "            .lower()\n",
        "            .replace(\"that feels\", \"feels\")\n",
        "            .replace(\"it feels\", \"feels\")\n",
        "            .replace(\"feels good\", \"feelsgood\")\n",
        "            .replace(\"good bye\", \"goodbye\")\n",
        "            .replace(\"good night\", \"goodnight\")\n",
        "            .replace(\"thank you\", \"thankyou\")\n",
        "            .replace(\"aaaaaa\", \"a\")\n",
        "            .replace(\"aaaa\", \"a\")\n",
        "            .replace(\"aa\", \"a\")\n",
        "            .replace(\"aa\", \"a\")\n",
        "            .replace(\"mmmmmm\", \"m\")\n",
        "            .replace(\"mmmm\", \"m\")\n",
        "            .replace(\"mm\", \"m\")\n",
        "            .replace(\"mm\", \"m\")\n",
        "            .replace(\"hhhhhh\", \"h\")\n",
        "            .replace(\"hhhh\", \"h\")\n",
        "            .replace(\"hh\", \"h\")\n",
        "            .replace(\"hh\", \"h\")\n",
        "            .replace(\"oooooo\", \"o\")\n",
        "            .replace(\"oooo\", \"o\")\n",
        "            .replace(\"oo\", \"o\")\n",
        "            .replace(\"oo\", \"o\")\n",
        "        )\n",
        "        is_garbage = True\n",
        "        for w in c.split(\" \"):\n",
        "            if w.strip() == \"\":\n",
        "                continue\n",
        "            if w.strip() in garbage_list:\n",
        "                continue\n",
        "            elif w.strip() in need_context_lines and last_line_garbage:\n",
        "                continue\n",
        "            else:\n",
        "                is_garbage = False\n",
        "                break\n",
        "        if not is_garbage:\n",
        "            clean_subs.append(subs[i])\n",
        "        last_line_garbage = is_garbage\n",
        "    with open(out_path, \"w\", encoding=\"utf8\") as f:\n",
        "        f.write(srt.compose(clean_subs))\n",
        "    print(\"\\nDone! Subs written to\", out_path)\n",
        "    print(\"Downloading SRT file:\")\n",
        "    files.download(out_path)\n"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}