whisperx_for_uploading_file.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# WhisperX for uploading files\n",
        "\n",
        "Upload your local files to the Colab Files from the left sidebar.\n",
        "\n",
        "从左侧将视频音频文件上传到Colab，然后运行即可下载生成好的字幕文件。\n",
        "\n"
      ],
      "metadata": {
        "id": "0Q0jEeQOvZmK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# @title *通用参数/Required settings:**\n",
        "# Check GPU availability\n",
        "!nvidia-smi\n",
        "\n",
        "# @markdown #### **Initial prompt**\n",
        "# @markdown Prompts can be very helpful for correcting specific words or acronyms that the model often misrecognizes in the audio.\n",
        "prompt = \"ChatGPT, LLM, Fine-tuning, Llama.\"  # @param {type:\"string\"}\n",
        "\n",
        "# @markdown #### **Directory Path**\n",
        "# @markdown where your audio-video files are located from Coloab\n",
        "directory_path = \"/content\"  # @param {type:\"string\"}\n",
        "\n",
        "# @markdown #### Model\n",
        "model_size = \"large-v2\"  # @param [\"base\", \"base.en\", \"small\", \"small.en\",\"medium\", \"medium.en\", \"large-v1\",\"large-v2\"]\n",
        "\n",
        "# @markdown #### Language\n",
        "language = \"auto\" # @param [\"auto\", \"en\", \"zh\", \"de\", \"es\", \"ru\", \"ko\", \"fr\", \"ja\", \"pt\", \"tr\", \"pl\", \"ca\", \"nl\", \"ar\", \"sv\", \"it\", \"id\", \"hi\", \"fi\", \"vi\", \"he\", \"uk\", \"el\", \"ms\", \"cs\", \"ro\", \"da\", \"hu\", \"ta\", \"no\", \"th\", \"ur\", \"hr\", \"bg\", \"lt\", \"la\", \"mi\", \"ml\", \"cy\", \"sk\", \"te\", \"fa\", \"lv\", \"bn\", \"sr\", \"az\", \"sl\", \"kn\", \"et\", \"mk\", \"br\", \"eu\", \"is\", \"hy\", \"ne\", \"mn\", \"bs\", \"kk\", \"sq\", \"sw\", \"gl\", \"mr\", \"pa\", \"si\", \"km\", \"sn\", \"yo\", \"so\", \"af\", \"oc\", \"ka\", \"be\", \"tg\", \"sd\", \"gu\", \"am\", \"yi\", \"lo\", \"uz\", \"fo\", \"ht\", \"ps\", \"tk\", \"nn\", \"mt\", \"sa\", \"lb\", \"my\", \"bo\", \"tl\", \"mg\", \"as\", \"tt\", \"haw\", \"ln\", \"ha\", \"ba\", \"jw\", \"su\"]\n",
        "\n",
        "# @markdown #### Filename Type\n",
        "# @markdown Use YouTube title as file name by default\n",
        "filename_type = \"title\"  # @param [\"title\", \"id\"]\n",
        "\n",
        "# @markdown #### Assign speaker labels\n",
        "# @markdown Recognize speakers\n",
        "assign_speaker_lable = False # @param {type:\"boolean\"}\n",
        "\n",
        "# @markdown #### Align whisper output\n",
        "align_whisper_output = True # @param {type:\"boolean\"}"
      ],
      "metadata": {
        "id": "J9Pgww7RuFUW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "80yOVjb-vAsD"
      },
      "source": [
        "# Install WhisperX"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Hcc35ui3ux8l"
      },
      "outputs": [],
      "source": [
        "! pip install git+https://github.com/m-bain/whisperx.git\n",
        "\n",
        "! pip install git+https://github.com/openai/whisper.git\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [],
      "metadata": {
        "id": "998j5EKSuFy7"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QxEqkpyuQL5_"
      },
      "source": [
        "# Get Subtitle Files using WhisperX"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qWk9Y3Uxv9qu"
      },
      "outputs": [],
      "source": [
        "\n",
        "import os\n",
        "from google.colab import files\n",
        "\n",
        "# supported extensions\n",
        "supported_extensions = ['.mp4', '.wav', '.mp3']\n",
        "\n",
        "language_param = \"\"\n",
        "if language != \"auto\":\n",
        "    language_param = f\"--language {language}\"\n",
        "\n",
        "diarize_param = \"\"\n",
        "if assign_speaker_lable:\n",
        "    diarize_param = \"--diarize --hf_token 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\"\n",
        "\n",
        "align_whisper_param = \"\"\n",
        "if align_whisper_output:\n",
        "    align_whisper_param = \"--align_model WAV2VEC2_ASR_LARGE_LV60K_960H\"\n",
        "\n",
        "prompt_param = \"\"\n",
        "if prompt != \"\":\n",
        "    prompt_param = f'--initial_prompt \"{prompt}\"'\n",
        "\n",
        "def process_file(filename):\n",
        "    # run = f'whisperx \"/content/APO2992689654.mp3\" --max_line_count 1 --max_line_width 100 --model medium.en --diarize --hf_token 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 --output_dir . --align_model WAV2VEC2_ASR_LARGE_LV60K_960H'\n",
        "    run = f'whisperx \"{filename}\" --model {model_size} {language_param} --output_dir . {prompt_param} {diarize_param} {align_whisper_param}'\n",
        "    !{run}\n",
        "\n",
        "    print(\"Start to download subtitle files\")\n",
        "    # start to download file\n",
        "    base_filename = os.path.splitext(filename)[0]\n",
        "    srt_filename = base_filename + '.srt'\n",
        "    json_filename = base_filename + '.json'\n",
        "\n",
        "    files.download(srt_filename)\n",
        "    files.download(json_filename)\n",
        "\n",
        "def process():\n",
        "    for root, dirs, files in os.walk(directory_path):\n",
        "        for file in files:\n",
        "            if file.endswith(tuple(supported_extensions)):\n",
        "                file_path = os.path.join(root, file)\n",
        "                process_file(file_path)\n",
        "\n",
        "process()"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}