N46Whisper.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cEAvBF4zgXX4"
      },
      "source": [
        "# N46Whisper\n",
        "\n",
        "N46Whisper is a Google Colab notebook application for streamlined video subtitle file generation.The original purpose of the project was to improve the productivity of Nogizaka46 (and Sakamichi groups) subbers. However, it can also be used to create subtitles in general.The application could significantly reduce the labour and time costs of sub-groups or individual subbers. However, despite its impressive performance, the Whisper model, AI translation and the application itself are not without limitations.\n",
        "\n",
        "\n",
        "N46Whisper 是基于 Google Colab 的应用。开发初衷旨在提高乃木坂46（以及坂道系）字幕组日语视频的制作效率,但亦适于所有外语视频的字幕制作。本应用的目标并非生产完美的字幕文件， 而旨在于搭建并提供一个简单且自动化的使用平台以节省生产成品字幕的时间和精力。Whisper模型有其本身的应用场景限制，AI 翻译的质量亦还不能尽如人意。\n",
        "\n",
        "<font size='4'>**对于中文用户，推荐在使用前阅读[常见问题说明](https://github.com/Ayanaminn/N46Whisper/blob/main/FAQ.md)。如果你觉得本应用对你有所帮助，欢迎帮助扩散给更多的人。**\n",
        "\n",
        "\n",
        "<font size='4'>**联系作者/Contact me：[E-mail](admin@ikedateresa.cc)**\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tZ2x8S7RMF9i"
      },
      "source": [
        "## 更新/What's Latest：\n",
        "历史更新日志\n",
        "\n",
        "<font size = '3'>**本项目将不再进行维护和更新，感谢大家的帮助与支持。**\n",
        "</br></br>\n",
        "\n",
        "2024.4.17:\n",
        "* 添加使用Google Gemini API翻译的选项。\n",
        "\n",
        "2024.1.31:\n",
        "* 鉴于集成的参数选项（还会）越来越多有使流程变得繁琐的趋势，这有违开发初衷。因此测试分离了一个[轻量版](https://colab.research.google.com/github/Ayanaminn/N46Whisper/blob/dev/N46WhisperLite.ipynb)，只保留最少的必要操作。\n",
        "\n",
        "2023.12.4:\n",
        "* 支持基于faster-whisper的WhisperV3模型/Support faster-whisper based WhisperV3 model\n",
        "\n",
        "2023.11.7:\n",
        "* 现在可以加载最新的WhisperV3模型/Enable users to load lastest Whisper V3 model.\n",
        "* 允许用户自行设置beam size/ Enable customerize beam size parameter.\n",
        "\n",
        "2023.4.30:\n",
        "* 优化提示词/Refine the translation prompt.\n",
        "* 允许用户使用个人提示词并调节Temperature参数/Allow user to custom prompt and temperature for translation.\n",
        "* 显示翻译任务消费统计/Display the token used and total cost for the translation task.\n",
        "\n",
        "2023.4.15:\n",
        "* 使用faster-whisper模型重新部署以提高效率，节省资源。Reimplement Whsiper based on faster-whsiper to improve efficiency.\n",
        "* 提供faster-whisper集成的vad filter选项以提高转录精度。Enable vad filter that integrated in faster-whisper to improve transcribe accuracy"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "laXaXrgPvCOn"
      },
      "source": [
        "**<font size='5'>以下选择文件方式按需执行其中一种即可，不需要全部运行</font>**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "30dM9s2J27A5"
      },
      "outputs": [],
      "source": [
        "#@title **从谷歌网盘选择文件/Select File From Google Drive**\n",
        "\n",
        "# @markdown <font size=\"2\">Navigate to the file you want to transcribe, left-click to highlight the file, then click 'Select' button to confirm.\n",
        "# @markdown <br/>从网盘目录中选择要转换的文件(视频/音频），单击选中文件，点击'Select'按钮以确认。</font><br/>\n",
        "# @markdown <br/><font size=\"2\">If use local file, ignore this cell and move to the next.\n",
        "# @markdown <br/>若希望从本地上传文件，则跳过此步执行下一单元格。</font><br/>\n",
        "# @markdown <br/><font size=\"2\">If file uploaded to drive after execution, execute this cell again to refresh.\n",
        "# @markdown <br/>若到这一步才上传文件到谷歌盘，则重复执行本单元格以刷新文件列表。</font>\n",
        "!pip install geemap\n",
        "from google.colab import drive\n",
        "from google.colab import files\n",
        "import os\n",
        "import logging\n",
        "from IPython.display import clear_output\n",
        "import geemap\n",
        "\n",
        "clear_output()\n",
        "drive.mount('/drive')\n",
        "\n",
        "print('Google Drive is mounted，please select file')\n",
        "print('谷歌云盘挂载完毕，请选择要转换的文件')\n",
        "\n",
        "from ipytree import Tree, Node\n",
        "import ipywidgets as widgets\n",
        "from ipywidgets import interactive\n",
        "# import os\n",
        "from google.colab import output\n",
        "output.enable_custom_widget_manager()\n",
        "use_drive = True\n",
        "global drive_dir\n",
        "drive_dir = []\n",
        "\n",
        "def file_tree():\n",
        "    # create widgets as a simple file browser\n",
        "    full_widget = widgets.HBox()\n",
        "    left_widget = widgets.VBox()\n",
        "    right_widget = widgets.VBox()\n",
        "\n",
        "    path_widget = widgets.Text()\n",
        "    path_widget.layout.min_width = '300px'\n",
        "    select_widget = widgets.Button(\n",
        "      description='Select', button_style='primary', tooltip='Select current media file.'\n",
        "      )\n",
        "    drive_url = widgets.Output()\n",
        "\n",
        "    right_widget.children = [select_widget]\n",
        "    full_widget.children = [left_widget]\n",
        "\n",
        "    tree_widget = widgets.Output()\n",
        "    tree_widget.layout.max_width = '300px'\n",
        "    tree_widget.overflow = 'auto'\n",
        "\n",
        "    left_widget.children = [path_widget,tree_widget]\n",
        "\n",
        "    # init file tree\n",
        "    my_tree = Tree(multiple_selection=False)\n",
        "    my_tree_dict = {}\n",
        "    media_names = []\n",
        "\n",
        "    def select_file(b):\n",
        "        drive_dir.append(path_widget.value)\n",
        "        # full_widget.disabled = True\n",
        "        # clear_output()\n",
        "        print('File selected，please continue to select more or execute next cell')\n",
        "        print('已选择文件，可以继续选择或执行下个单元格')\n",
        "    #     if (out_file not in my_tree_dict.keys()) and (out_dir in my_tree_dict.keys()):\n",
        "    #         node = Node(os.path.basename(out_file))\n",
        "    #         my_tree_dict[out_file] = node\n",
        "    #         parent_node = my_tree_dict[out_dir]\n",
        "    #         parent_node.add_node(node)\n",
        "\n",
        "    select_widget.on_click(select_file)\n",
        "\n",
        "    def handle_file_click(event):\n",
        "        if event['new']:\n",
        "            cur_node = event['owner']\n",
        "            for key in my_tree_dict.keys():\n",
        "                if (cur_node is my_tree_dict[key]) and (os.path.isfile(key)):\n",
        "                    try:\n",
        "                        with open(key) as f:\n",
        "                            path_widget.value = key\n",
        "                            path_widget.disabled = False\n",
        "                            select_widget.disabled = False\n",
        "                            full_widget.children = [left_widget, right_widget]\n",
        "                    except Exception as e:\n",
        "                        path_widget.value = key\n",
        "                        path_widget.disabled = True\n",
        "                        select_widget.disabled = True\n",
        "\n",
        "                        return\n",
        "\n",
        "    def handle_folder_click(event):\n",
        "        if event['new']:\n",
        "            full_widget.children = [left_widget]\n",
        "\n",
        "    # redirect cwd to default drive root path and add nodes\n",
        "    my_dir = '/drive/MyDrive'\n",
        "    my_root_name = my_dir.split('/')[-1]\n",
        "    my_root_node = Node(my_root_name)\n",
        "    my_tree_dict[my_dir] = my_root_node\n",
        "    my_tree.add_node(my_root_node)\n",
        "    my_root_node.observe(handle_folder_click, 'selected')\n",
        "\n",
        "    for root, d_names, f_names in os.walk(my_dir):\n",
        "        folders = root.split('/')\n",
        "        for folder in folders:\n",
        "            if folder.startswith('.'):\n",
        "                continue\n",
        "        for d_name in d_names:\n",
        "            if d_name.startswith('.'):\n",
        "                d_names.remove(d_name)\n",
        "        for f_name in f_names:\n",
        "            # if f_name.startswith('.'):\n",
        "            #     f_names.remove(f_name)\n",
        "            # only add media files\n",
        "            if f_name.lower().endswith(('mp3','m4a','flac','aac','wav','mp4','mkv','ts','flv')):\n",
        "                media_names.append(f_name)\n",
        "\n",
        "        d_names.sort()\n",
        "        f_names.sort()\n",
        "        media_names.sort()\n",
        "        keys = my_tree_dict.keys()\n",
        "\n",
        "        if root not in my_tree_dict.keys():\n",
        "          # print(f'root name is {root}') # folder path\n",
        "          name = root.split('/')[-1] # folder name\n",
        "          # print(f'folder name is {name}')\n",
        "          dir_name = os.path.dirname(root) # parent path of folder\n",
        "          # print(f'dir name is {dir_name}')\n",
        "          parent_node = my_tree_dict[dir_name]\n",
        "          node = Node(name)\n",
        "          my_tree_dict[root] = node\n",
        "          parent_node.add_node(node)\n",
        "          node.observe(handle_folder_click, 'selected')\n",
        "\n",
        "        if len(media_names) > 0:\n",
        "              parent_node = my_tree_dict[root] # parent folders\n",
        "              # print(parent_node)\n",
        "              parent_node.opened = False\n",
        "              for f_name in media_names:\n",
        "                  node = Node(f_name)\n",
        "                  node.icon = 'file'\n",
        "                  full_path = os.path.join(root, f_name)\n",
        "                  # print(full_path)\n",
        "                  my_tree_dict[full_path] = node\n",
        "                  parent_node.add_node(node)\n",
        "                  node.observe(handle_file_click, 'selected')\n",
        "        media_names.clear()\n",
        "\n",
        "    with tree_widget:\n",
        "      tree_widget.clear_output()\n",
        "      display(my_tree)\n",
        "\n",
        "    return full_widget\n",
        "\n",
        "\n",
        "tree= file_tree()\n",
        "tree\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "VF_eOkuYbrBj"
      },
      "outputs": [],
      "source": [
        "#@title **从本地上传文件(可多选）/Upload Local File（Can select multiple)**\n",
        "# @markdown <font size=\"2\">If use file in google drive, ignore this cell and move to the next.\n",
        "# @markdown <br/>若已选择谷歌盘中的文件，则跳过此步执行下一单元格。</font>\n",
        "\n",
        "from google.colab import files\n",
        "use_drive = False\n",
        "uploaded = files.upload()\n",
        "file_names = []\n",
        "file_names.append(list(uploaded.keys())[0])\n",
        "print('File uploaded，please continue to upload more or execute next cell')\n",
        "print('已上传文件，可以执行下个单元格')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y8kquVjWvh6c"
      },
      "source": [
        "**<font size='5'>以下顺次点击下方每个单元格左侧的“运行”图标，不可跳过步骤</font>**\n",
        "**</br>【重要】:** 务必在\"修改\"->\"笔记本设置\"->\"硬件加速器\"中选择GPU！否则处理速度会非常慢。\n",
        " **</br>【IMPORTANT】:** Make sure you select GPU as hardware accelerator in notebook settings, otherwise the processing speed will be very slow."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "gLcsoJy5BIcW"
      },
      "outputs": [],
      "source": [
        "#@title **通用参数/Required settings:**\n",
        "\n",
        "\n",
        "# @markdown **【IMPORTANT】:**<font size=\"2\">Select uploaded file type.\n",
        "# @markdown **</br>【重要】:** 选择上传的文件类型(视频-video/音频-audio）</font>\n",
        "\n",
        "# encoding:utf-8\n",
        "file_type = \"audio\"  # @param [\"audio\",\"video\"]\n",
        "\n",
        "# @markdown <font size=\"2\">Model size will affect the processing time and transcribe quality.\n",
        "# @markdown <br/>The default source language is Japanese.Please input your own source language if applicable.Use two letter language code， e.g.  'en', 'ja'...\n",
        "# @markdown <br/>模型大小将影响转录时间和质量, **默认使用稳定的large-v2模型以节省时间**\n",
        "# @markdown <br/>默认识别语言为日语，若使用其它语言的视频请自行输入即可。请注意：使用两字母语言代码如'en'，'ja'\n",
        "# @markdown <br/>【请注意】：large-v3在某些情况下可能未必优于large-v2或更早的模型，请用户自行选择\n",
        "\n",
        "model_size = \"large-v2\"  # @param [\"base\",\"small\",\"medium\", \"large-v1\",\"large-v2\",\"large-v3\"]\n",
        "language = \"ja\"  # @param {type:\"string\"}\n",
        "\n",
        "# @markdown <font size=\"2\">默认只导出ass，若需要srt则选择Yes</font>\n",
        "# @markdown <br/><font size=\"2\">导出时浏览器会弹出允许同时下载多个文件的请求，需要同意\n",
        "export_srt = \"No\"  # @param [\"No\",\"Yes\"]\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "nw72_bK3AS1d"
      },
      "outputs": [],
      "source": [
        "#@title **其他选项/Advanced settings**\n",
        "\n",
        "# @markdown <font size=\"2\">Option for split line text by spaces. The splited lines all use the same time stamp, with 'adjust_required' label as remark for manual adjustment.\n",
        "# @markdown <br/>将存在空格的单行文本分割为多行（多句）。分割后的若干行均临时采用相同时间戳，且添加了adjust_required标记提示调整时间戳避免叠轴\n",
        "# @markdown <br/>普通分割（Modest): 当空格后的文本长度超过5个字符，则另起一行\n",
        "# @markdown <br/>全部分割（Aggressive): 只要遇到空格即另起一行\n",
        "# @markdown <br/>标点分割（Punctuation): 只要遇到句号即另起一行，在未来可能添加更加智能的标点分割方法\n",
        "is_split = \"No\"  # @param [\"No\",\"Yes\"]\n",
        "split_method = \"Modest\"  # @param [\"Modest\",\"Aggressive\", \"Punctuation\"]\n",
        "# @markdown <font size=\"2\">Please contact us if you want to have your sub style integrated.\n",
        "# @markdown <br/>当前支持生成字幕格式：\n",
        "# @markdown <br/><li>ikedaCN - 特蕾纱熊猫观察会字幕组\n",
        "# @markdown <br/><li>sugawaraCN - 坂上之月字幕组\n",
        "# @markdown <br/><li>kaedeCN - 三番目の枫字幕组\n",
        "# @markdown <br/><li>taniguchiCN - 泪痣愛季応援団\n",
        "# @markdown <br/><li>asukaCN - 暗鳥其实很甜字幕组\n",
        "sub_style = \"default\"  # @param [\"default\", \"ikedaCN\", \"kaedeCN\",\"sugawaraCN\",\"taniguchiCN\",\"asukaCN\"]\n",
        "\n",
        "# @markdown **使用VAD过滤/Use VAD filter**\n",
        "\n",
        "# @markdown <font size=\"2\">使用[Silero VAD model](https://github.com/snakers4/silero-vad)以检测并过滤音频中的无声段落（推荐小语种使用）\n",
        "# @markdown <br/>[WARNING] Use VAD filter have pros and cons, please carefully select this option accroding to your own audio content.\n",
        "# @markdown <br/>【注意】使用VAD filter有优点亦有缺点，请用户自行根据音频内容决定是否启用. [关于VAD filter](https://github.com/Ayanaminn/N46Whisper/blob/main/FAQ.md)\n",
        "\n",
        "\n",
        "is_vad_filter = \"False\" # @param [\"True\", \"False\"]\n",
        "# @markdown  <font size=\"2\"> *  The default <font size=\"3\">  ```min_silence_duration``` <font size=\"2\"> is set at 1000 ms in N46Whisper\n",
        "\n",
        "# @markdown **设置Beam Size**\n",
        "\n",
        "# @markdown <font size=\"2\">Beam Size数值越高，在识别时探索的路径越多，这在一定范围内可以帮助提高识别准确性，但是相对的VRAM使用也会更高. 同时，Beam Size在超过5-10后有可能降低精确性，详情请见https://arxiv.org/pdf/2204.05424.pdf\n",
        "# @markdown <br/> 默认设置为 5\n",
        "set_beam_size = 5 #@param\n",
        "\n",
        "# @markdown <font size=\"2\">在不设置Beam Size时，Whisper将会使用贪心解码，这在一定程度上可能与英语等其他语言的换行功能有联系，详情请见https://github.com/Ayanaminn/N46Whisper/issues/46\n",
        "# @markdown <br/> 默认设置为 false\n",
        "beam_size_off = False # @param {type:\"boolean\"}\n",
        "\n",
        "# # @markdown <font size=\"2\">设置此参数为True将在代码执行完毕后自动断开Colab的连接。这有助于在长时间运行的任务完成后释放资源。请注意，断开连接后，所有未保存的数据将丢失。</font>\n",
        "# # @markdown <br/> 默认设置为 False\n",
        "# auto_disconnect = True #@param {type:\"boolean\"}\n",
        "\n",
        "# # 以下代码用于断开连接\n",
        "# from IPython.display import Javascript\n",
        "\n",
        "# def disconnect_runtime():\n",
        "#     if auto_disconnect:\n",
        "#         display(Javascript('google.colab.kernel.disconnect();'))\n",
        "#         print(\"已经自动断开连接。\")\n",
        "#     else:\n",
        "#         print(\"自动断开连接功能已关闭。\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "z0igG7ruI-7q"
      },
      "outputs": [],
      "source": [
        "#@title **运行Whisper/Run Whisper**\n",
        "# Hugging Face Hub\n",
        "# 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\n",
        "\n",
        "#@markdown 完成后ass文件将自动下载到本地/ass file will be auto downloaded after finish.\n",
        "! pip install ffmpeg\n",
        "! wget https://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@raw.githubusercontent.com/Ayanaminn/N46Whisper/main/srt2ass.py\n",
        "! pip install pysubs2\n",
        "! pip install faster-whisper\n",
        "# ! pip install nest_asyncio\n",
        "import torch\n",
        "# import nest_asyncio\n",
        "from faster_whisper import WhisperModel\n",
        "import ipywidgets as widgets\n",
        "from IPython.display import display, clear_output\n",
        "clear_output()\n",
        "print('语音识别库配置完毕，将开始转换')\n",
        "import os\n",
        "import ffmpeg\n",
        "import subprocess\n",
        "from tqdm import tqdm\n",
        "import time\n",
        "import pandas as pd\n",
        "import requests\n",
        "from urllib.parse import quote_plus\n",
        "from pathlib import Path\n",
        "import sys\n",
        "import pysubs2\n",
        "import gc\n",
        "import zipfile\n",
        "import asyncio\n",
        "\n",
        "# Enable nested asyncio in Jupyter Notebook\n",
        "# nest_asyncio.apply()\n",
        "\n",
        "# assert file_name != \"\"\n",
        "# assert language != \"\"\n",
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\")\n",
        "\n",
        "file_basenames = []\n",
        "\n",
        "if use_drive:\n",
        "    output_dir = os.path.dirname(drive_dir[0])\n",
        "    try:\n",
        "        file_names = drive_dir\n",
        "        for i in range(len(file_names)):\n",
        "          file_basenames.append(file_names[i].split('.')[0])\n",
        "        # print(file_name)\n",
        "        output_dir = os.path.dirname(drive_dir[0])\n",
        "    except Exception as e:\n",
        "            print(f'error: {e}')\n",
        "else:\n",
        "    sys.path.append('/drive/content')\n",
        "    if not os.path.exists(file_names[0]):\n",
        "      raise ValueError(f\"No {file_names[0]} found in current path.\")\n",
        "    else:\n",
        "        try:\n",
        "            for i in range(len(file_names)):\n",
        "              file_basenames.append(Path(file_names[i]).stem)\n",
        "            output_dir = Path(file_names[0]).parent.resolve()\n",
        "            # print(file_basename)\n",
        "            # print(output_dir)\n",
        "        except Exception as e:\n",
        "            print(f'error: {e}')\n",
        "\n",
        "\n",
        "clear_output()\n",
        "print('加载模型 Loading model...')\n",
        "\n",
        "model = WhisperModel(model_size)\n",
        "torch.cuda.empty_cache()\n",
        "\n",
        "for i in range(len(file_names)):\n",
        "  file_name = file_names[i]\n",
        "  #Transcribe\n",
        "  file_basename = file_basenames[i]\n",
        "  if file_type == \"video\":\n",
        "    print('提取音频中 Extracting audio from video file...')\n",
        "    os.system(f'ffmpeg -i {file_name} -f mp3 -ab 192000 -vn {file_basename}.mp3')\n",
        "    print('提取完毕 Done.')\n",
        "  # print(file_basename)\n",
        "  tic = time.time()\n",
        "  clear_output()\n",
        "  print('识别中 Transcribe in progress...')\n",
        "\n",
        "  if beam_size_off:\n",
        "    segments, info = model.transcribe(audio = f'{file_name}',\n",
        "                                          language=language,\n",
        "                                          vad_filter=is_vad_filter,\n",
        "                                          vad_parameters=dict(min_silence_duration_ms=1000))\n",
        "  else:\n",
        "    segments, info = model.transcribe(audio = f'{file_name}',\n",
        "                                          beam_size=set_beam_size,\n",
        "                                          language=language,\n",
        "                                          vad_filter=is_vad_filter,\n",
        "                                          vad_parameters=dict(min_silence_duration_ms=1000))\n",
        "\n",
        "  # segments is a generator so the transcription only starts when you iterate over it\n",
        "  # to use pysubs2, the argument must be a segment list-of-dicts\n",
        "  total_duration = round(info.duration, 2)  # Same precision as the Whisper timestamps.\n",
        "  results= []\n",
        "  with tqdm(total=total_duration, unit=\" seconds\") as pbar:\n",
        "      for s in segments:\n",
        "          segment_dict = {'start':s.start,'end':s.end,'text':s.text}\n",
        "          results.append(segment_dict)\n",
        "          segment_duration = s.end - s.start\n",
        "          pbar.update(segment_duration)\n",
        "\n",
        "\n",
        "  #Time comsumed\n",
        "  toc = time.time()\n",
        "  print('识别完毕 Done')\n",
        "  print(f'Time consumpution {toc-tic}s')\n",
        "\n",
        "  subs = pysubs2.load_from_whisper(results)\n",
        "  srt_filename = file_basename + '.srt'\n",
        "  subs.save(srt_filename)\n",
        "\n",
        "  from srt2ass import srt2ass\n",
        "  ass_filename  = srt2ass(srt_filename, sub_style, is_split,split_method)\n",
        "  print('ASS subtitle saved as: ' + ass_filename )\n",
        "\n",
        "  if i+1 ==1:\n",
        "    #一个文件则直接下载\n",
        "    files.download(ass_filename )\n",
        "  else:\n",
        "    # 用压缩包来处理所有下载，避免多个下载触发浏览器限制，必须要跟用户交互才能继续\n",
        "    zip_filename = file_basename+\".zip\"\n",
        "    with zipfile.ZipFile(zip_filename, 'w') as zipf:\n",
        "      if export_srt == 'Yes':\n",
        "          zipf.write(srt_filename)\n",
        "      zipf.write(ass_filename)\n",
        "\n",
        "    # Trigger the download asynchronously\n",
        "    files.download(zip_filename)\n",
        "\n",
        "  print('第',i+1,'个文件字幕生成完毕/',i+1, 'file(s) was completed!')\n",
        "  torch.cuda.empty_cache()\n",
        "\n",
        "print('所有字幕生成完毕 All done!')\n",
        "\n",
        "# 如果设定了自动关闭连接，这里将释放Colab资源\n",
        "# disconnect_runtime()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "k5n2xrB631JV"
      },
      "outputs": [],
      "source": [
        "#@title **【实验功能】Experimental Features:**\n",
        "\n",
        "# @markdown **AI文本翻译/AI Translation:**\n",
        "# @markdown **</br>**<font size=\"2\"> 此功能允许用户使用AI翻译服务对识别的字幕文件做逐行翻译，并以相同的格式生成双语对照字幕。\n",
        "# @markdown **</br>**阅读项目文档以了解更多。</font>\n",
        "# @markdown **</br>**<font size=\"2\"> This feature allow users to translate previously transcribed subtitle text line by line using AI translation.\n",
        "# @markdown **</br>**Then generate bilingual subtitle files in same sub style.Read documentaion to learn more.</font>\n",
        "\n",
        "# @markdown **</br>**希望在本地使用字幕翻译功能的用户，推荐尝试 [subtitle-translator-electron](https://github.com/gnehs/subtitle-translator-electron)\n",
        "\n",
        "# @markdown **</br><font size=\"3\">Select subtitle file source</br>\n",
        "# @markdown <font size=\"3\">选择字幕文件(使用上一步的转录-use_transcribed/新上传-upload_new）</br>**\n",
        "# @markdown <font size=\"2\">支持SRT与ASS文件\n",
        "sub_source = \"upload_new\"  # @param [\"use_transcribed\",\"upload_new\"]\n",
        "\n",
        "# @markdown **chatGPT:**\n",
        "# @markdown **</br>**<font size=\"2\"> 要使用chatGPT翻译，请填入你自己的OpenAI API Key，目标语言，输出类型，然后执行单元格。</font>\n",
        "# @markdown **</br>**<font size=\"2\"> Please input your own OpenAI API Key, then execute this cell.</font>\n",
        "# @markdown **</br>**<font size=\"2\">【注意】 免费的API对速度有所限制，需要较长时间，用户可以自行考虑付费方案。</font>\n",
        "# @markdown **</br>**<font size=\"2\">【Note】There are limitaions on usage for free API, consider paid plan to speed up.</font>\n",
        "openai_key = '' # @param {type:\"string\"}\n",
        "target_language = 'zh-hans'# @param [\"zh-hans\",\"english\"]\n",
        "prompt = \"You are a language expert.Your task is to translate the input subtitle text, sentence by sentence, into the user specified target language.However, please utilize the context to improve the accuracy and quality of translation.Please be aware that the input text could contain typos and grammar mistakes, utilize the context to correct the translation.Please return only translated content and do not include the origin text.Please do not use any punctuation around the returned text.Please do not translate people's name and leave it as original language.\\\"\" # @param {type:\"string\"}\n",
        "temperature = 0.6 #@param {type:\"slider\", min:0, max:1.0, step:0.1}\n",
        "# @markdown <font size=\"4\">Default prompt: </br>\n",
        "# @markdown ```You are a language expert.```</br>\n",
        "# @markdown ```Your task is to translate the input subtitle text, sentence by sentence, into the user specified target language.```</br>\n",
        "# @markdown ```Please utilize the context to improve the accuracy and quality of translation.```</br>\n",
        "# @markdown ```Please be aware that the input text could contain typos and grammar mistakes, utilize the context to correct the translation.```</br>\n",
        "# @markdown ```Please return only translated content and do not include the origin text.```</br>\n",
        "# @markdown ```Please do not use any punctuation around the returned text.```</br>\n",
        "# @markdown ```Please do not translate people's name and leave it as original language.```</br>\n",
        "output_format = \"ass\"  # @param [\"ass\",\"srt\"]\n",
        "\n",
        "import sys\n",
        "import os\n",
        "import re\n",
        "import time\n",
        "import codecs\n",
        "import regex as re\n",
        "from pathlib import Path\n",
        "from tqdm import tqdm\n",
        "from google.colab import files\n",
        "from IPython.display import clear_output\n",
        "\n",
        "clear_output()\n",
        "\n",
        "if sub_source == 'upload_new':\n",
        "  uploaded = files.upload()\n",
        "  sub_name = list(uploaded.keys())[0]\n",
        "  sub_basename = Path(sub_name).stem\n",
        "elif sub_source == 'use_transcribed':\n",
        "  sub_name = file_basenames[0] +'.ass'\n",
        "  sub_basename = file_basenames[0]\n",
        "\n",
        "!pip install openai\n",
        "!pip install pysubs2\n",
        "from openai import OpenAI\n",
        "import pysubs2\n",
        "\n",
        "clear_output()\n",
        "\n",
        "class ChatGPTAPI():\n",
        "    def __init__(self, key, language, prompt, temperature):\n",
        "        self.key = key\n",
        "        # self.keys = itertools.cycle(key.split(\",\"))\n",
        "        self.language = language\n",
        "        self.key_len = len(key.split(\",\"))\n",
        "        self. prompt = prompt\n",
        "        self.temperature = temperature\n",
        "\n",
        "\n",
        "    # def rotate_key(self):\n",
        "    #     openai.api_key = next(self.keys)\n",
        "\n",
        "    def translate(self, text):\n",
        "        # print(text)\n",
        "        # self.rotate_key()\n",
        "        client = OpenAI(\n",
        "            api_key=self.key,\n",
        "            )\n",
        "\n",
        "        try:\n",
        "            completion = client.chat.completions.create(\n",
        "                model=\"gpt-3.5-turbo\",\n",
        "                messages=[\n",
        "                    {\n",
        "                        \"role\": \"system\",\n",
        "                        # english prompt here to save tokens\n",
        "                        \"content\": f'{self.prompt}'\n",
        "                    },\n",
        "                    {\n",
        "                        \"role\":\"user\",\n",
        "                        \"content\": f\"Original text:`{text}`. Target language: {self.language}\"\n",
        "                    }\n",
        "                ],\n",
        "                temperature=self.temperature\n",
        "            )\n",
        "            t_text = (\n",
        "                completion.choices[0].message.content.encode(\"utf8\").decode()\n",
        "            )\n",
        "            total_tokens = completion.usage.total_tokens # include prompt_tokens and completion_tokens\n",
        "        except Exception as e:\n",
        "            # TIME LIMIT for open api , pay to reduce the waiting time\n",
        "            sleep_time = int(60 / self.key_len)\n",
        "            time.sleep(sleep_time)\n",
        "            print(e, f\"will sleep  {sleep_time} seconds\")\n",
        "            # self.rotate_key()\n",
        "            client = OpenAI(\n",
        "            api_key=self.key,\n",
        "            )\n",
        "            completion = client.chat.completions.create(\n",
        "                model=\"gpt-3.5-turbo\",\n",
        "                messages=[\n",
        "                    {\n",
        "                        \"role\": \"system\",\n",
        "                        \"content\": f'{self.prompt}'\n",
        "                    },\n",
        "                    {\n",
        "                        \"role\": \"user\",\n",
        "                        \"content\": f\"Original text:`{text}`. Target language: {self.language}\"\n",
        "                    }\n",
        "                ],\n",
        "                temperature=self.temperature\n",
        "            )\n",
        "            t_text = (\n",
        "                completion.choices[0].message.content.encode(\"utf8\").decode()\n",
        "            )\n",
        "        total_tokens = completion.usage.total_tokens\n",
        "        return t_text, total_tokens\n",
        "\n",
        "\n",
        "class SubtitleTranslator():\n",
        "    def __init__(self, sub_src, model, key, language, prompt,temperature):\n",
        "        self.sub_src = sub_src\n",
        "        self.translate_model = model(key, language,prompt,temperature)\n",
        "        self.translations = []\n",
        "        self.total_tokens = 0\n",
        "\n",
        "    def calculate_price(self,num_tokens):\n",
        "        price_per_token = 0.000002 #gpt-3.5-turbo\t$0.002 / 1K tokens\n",
        "        return num_tokens * price_per_token\n",
        "\n",
        "    def translate_by_line(self):\n",
        "        sub_trans = pysubs2.load(self.sub_src)\n",
        "        total_lines = len(sub_trans)\n",
        "        for line in tqdm(sub_trans,total = total_lines):\n",
        "            line_trans, tokens_per_task = self.translate_model.translate(line.text)\n",
        "            line.text += (r'\\N'+ line_trans)\n",
        "            print(line_trans)\n",
        "            self.translations.append(line_trans)\n",
        "            self.total_tokens += tokens_per_task\n",
        "\n",
        "        return sub_trans, self.translations, self.total_tokens\n",
        "\n",
        "\n",
        "clear_output()\n",
        "\n",
        "translate_model = ChatGPTAPI\n",
        "\n",
        "assert translate_model is not None, \"unsupported model\"\n",
        "OPENAI_API_KEY = openai_key\n",
        "\n",
        "if not OPENAI_API_KEY:\n",
        "    raise Exception(\n",
        "        \"OpenAI API key not provided, please google how to obtain it\"\n",
        "    )\n",
        "# else:\n",
        "#     OPENAI_API_KEY = openai_key\n",
        "\n",
        "t = SubtitleTranslator(\n",
        "    sub_src=sub_name,\n",
        "    model= translate_model,\n",
        "    key = OPENAI_API_KEY,\n",
        "    language=target_language,\n",
        "    prompt=prompt,\n",
        "    temperature=temperature)\n",
        "\n",
        "translation, _, total_token = t.translate_by_line()\n",
        "total_price = t.calculate_price(total_token)\n",
        "#Download ass file\n",
        "\n",
        "if output_format == 'ass':\n",
        "  translation.save(sub_basename + '_translation.ass')\n",
        "  files.download(sub_basename + '_translation.ass')\n",
        "elif output_format == 'srt':\n",
        "  translation.save(sub_basename + '_translation.srt')\n",
        "  files.download(sub_basename + '_translation.srt')\n",
        "\n",
        "\n",
        "\n",
        "print('双语字幕生成完毕 All done!')\n",
        "print(f\"Total number of tokens used: {total_token}\")\n",
        "print(f\"Total price (USD): ${total_price:.4f}\")\n",
        "\n",
        "# @markdown **</br>**<font size='3'>**实验功能的开发亦是为了尝试帮助大家更有效率的制作字幕。但是只有在用户实际使用体验反馈的基础上，此应用才能不断完善，如果您有任何想法，都欢迎以任何方式联系我，提出[issue](https://github.com/Ayanaminn/N46Whisper/issues)或者分享在[讨论区](https://github.com/Ayanaminn/N46Whisper/discussions)。**\n",
        "# @markdown **</br>**<font size='3'>**The efficacy of this application cannot get improved without the feedbacks from everyday users.Please feel free to share your thoughts with me or post it [here](https://github.com/Ayanaminn/N46Whisper/discussions)**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "52JugZiQMQUm"
      },
      "outputs": [],
      "source": [
        "#@title **【实验功能】Experimental Features:**\n",
        "\n",
        "# @markdown **Google Gemini AI文本翻译/Google Gemini AI Translation:**\n",
        "# @markdown **</br>**<font size=\"2\"> 由于谷歌Gemini提供免费的API，想要免费使用AI翻译的用户可以执行该单元格。\n",
        "# @markdown **</br>**阅读项目文档以了解更多。</font>\n",
        "# @markdown **</br>**<font size=\"2\"> Since Google Gemini provides a free tier API, the users that want to use free AI translations can execute this block of code.\n",
        "# @markdown **</br>**Then generate bilingual subtitle files in same sub style.Read documentaion to learn more.</font>\n",
        "\n",
        "# @markdown **注意：同时执行Whisper翻译和Gemini API翻译有可能遇到runtime问题，建议在disconnect runtime之后重新执行该单元格**\n",
        "\n",
        "# @markdown **Attention: Executing Whisper and Gemini translation at the same time might run into runtime errors, we recommend that the users disconnect their runtimes before executing this block**\n",
        "output_format = \"ass\"  # @param [\"ass\",\"srt\"]\n",
        "\n",
        "# @markdown **Google Gemini:**\n",
        "# @markdown **</br>**<font size=\"2\"> 要使用Gemini翻译，请填入你自己的Gemini API Key，目标语言，输出类型，然后执行单元格。</font>\n",
        "# @markdown **</br>**<font size=\"2\"> Please input your own Gemini API Key, then execute this cell.</font>\n",
        "\n",
        "google_api_key = '' # @param {type:\"string\"}\n",
        "target_language = 'zh-hans'# @param [\"zh-hans\",\"english\"]\n",
        "prompt = \"You are a language expert.Your task is to translate the input subtitle text, sentence by sentence, into the user specified target language.However, please utilize the context to improve the accuracy and quality of translation.Please be aware that the input text could contain typos and grammar mistakes, utilize the context to correct the translation.Please return only translated content and do not include the origin text.Please do not use any punctuation around the returned text.Please do not translate people's name and leave it as original language.\\\"\" # @param {type:\"string\"}\n",
        "temperature = 0.6 #@param {type:\"slider\", min:0, max:1.0, step:0.1}\n",
        "# @markdown <font size=\"4\">Default prompt: </br>\n",
        "# @markdown ```You are a language expert.```</br>\n",
        "# @markdown ```Your task is to translate the input subtitle text, sentence by sentence, into the user specified target language.```</br>\n",
        "# @markdown ```Please utilize the context to improve the accuracy and quality of translation.```</br>\n",
        "# @markdown ```Please be aware that the input text could contain typos and grammar mistakes, utilize the context to correct the translation.```</br>\n",
        "# @markdown ```Please return only translated content and do not include the origin text.```</br>\n",
        "# @markdown ```Please do not use any punctuation around the returned text.```</br>\n",
        "# @markdown ```Please do not translate people's name and leave it as original language.```</br>\n",
        "\n",
        "import sys\n",
        "import os\n",
        "import re\n",
        "import time\n",
        "import codecs\n",
        "import regex as re\n",
        "from pathlib import Path\n",
        "from tqdm import tqdm\n",
        "from google.colab import files\n",
        "from IPython.display import clear_output\n",
        "\n",
        "uploaded = files.upload()\n",
        "sub_name = list(uploaded.keys())[0]\n",
        "sub_basename = Path(sub_name).stem\n",
        "\n",
        "clear_output()\n",
        "\n",
        "!pip install -q -U google-generativeai\n",
        "!pip install pysubs2\n",
        "\n",
        "import google.generativeai as genai\n",
        "from google.generativeai.types import HarmCategory, HarmBlockThreshold\n",
        "import pysubs2\n",
        "\n",
        "genai.configure(api_key=google_api_key)\n",
        "\n",
        "model = genai.GenerativeModel('gemini-pro')\n",
        "\n",
        "def translate(prompt, language, text, retry_times=0):\n",
        "        # print(text)\n",
        "        # self.rotate_key()\n",
        "        try:\n",
        "            completion = model.generate_content(prompt + \"target language :\" + language + \"text: \" + text, safety_settings={\n",
        "                    HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,\n",
        "                    HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,\n",
        "                    HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,\n",
        "                    HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,\n",
        "            })\n",
        "            t_text = (\n",
        "                completion.text\n",
        "            )\n",
        "            print(t_text)\n",
        "            return t_text\n",
        "        except Exception as e:\n",
        "            if re.search(r\"^429 POST|Remote end closed connection|Connection reset by peer\", str(e), flags=re.I) and retry_times < 6:\n",
        "                print(\"翻译接口请求过于频繁或被断开，将再次重试，重试次数: %d\" % (retry_times + 1))\n",
        "                print(\"The translation interface request is too frequent or disconnected. Will try again. Number of retries: %d\" % (retry_times + 1))\n",
        "                time.sleep(min(retry_times + 1, 3))\n",
        "                return translate(prompt, language, text, retry_times + 1)\n",
        "            else:\n",
        "                # Since Google API should not run into runtime error, this would be an unknown error\n",
        "                print(\"未知错误，用户可以尝试查看报错信息并在Repository里提交issue\")\n",
        "                print(\"Unknown Error, please check the error log and open an issue in the repository\")\n",
        "                return \">>>>> UnknownError\"\n",
        "class SubtitleTranslator():\n",
        "    def __init__(self, sub_src):\n",
        "        self.sub_src = sub_src\n",
        "        self.translations = []\n",
        "\n",
        "    def translate_by_line(self):\n",
        "        sub_trans = pysubs2.load(self.sub_src)\n",
        "        total_lines = len(sub_trans)\n",
        "        for line in tqdm(sub_trans,total = total_lines):\n",
        "            line_trans = translate(prompt, target_language, line.text)\n",
        "            line.text += (r'\\N'+ line_trans)\n",
        "            print(line_trans)\n",
        "            self.translations.append(line_trans)\n",
        "        return sub_trans, self.translations\n",
        "\n",
        "clear_output()\n",
        "\n",
        "if not google_api_key:\n",
        "    raise Exception(\n",
        "        \"Google Gemini API key not provided, please google how to obtain it\"\n",
        "    )\n",
        "\n",
        "t = SubtitleTranslator(sub_src=sub_name)\n",
        "\n",
        "translation, _, = t.translate_by_line()\n",
        "\n",
        "if output_format == 'ass':\n",
        "  translation.save(sub_basename + '_translation.ass')\n",
        "  files.download(sub_basename + '_translation.ass')\n",
        "elif output_format == 'srt':\n",
        "  translation.save(sub_basename + '_translation.srt')\n",
        "  files.download(sub_basename + '_translation.srt')\n",
        "\n",
        "print('双语字幕生成完毕 All done!')"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}