src/run_whisper_colab.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PcoXHrM6jGof"
      },
      "source": [
        "# Run WhisperAI using Google Colab's GPU\n",
        "\n",
        "To enable GPU in this notebook, select Runtime -> Change runtime type in the Menu bar. Under Hardware Accelerator, select GPU.\n",
        "\n",
        "Then, go to the Global variables [cell](#scrollTo=nEyzcieyHhav&line=3&uniqifier=1) and update it with your ngrok authentication token.\n",
        "\n",
        "To run WhisperAI, select Runtime -> Run all. Go to this [cell](#scrollTo=hCrK1VjzOpJk&line=8&uniqifier=1) and read the instructions on how to update your `.env` file to use WhisperAI running in this notebook."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LFOyJ9GkkoW-"
      },
      "source": [
        "## Pip install dependencies"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "STstw_cvPCjL"
      },
      "outputs": [],
      "source": [
        "!pip install flask-ngrok2 -q\n",
        "!pip install git+https://github.com/openai/whisper.git -q"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FVRHvVNmktDd"
      },
      "source": [
        "## Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r67_m4DoPN-X"
      },
      "outputs": [],
      "source": [
        "from flask_ngrok2 import run_with_ngrok\n",
        "from flask import Flask, request \n",
        "import json\n",
        "import whisper"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G4xDe-tKkvr4"
      },
      "source": [
        "## Global variables\n",
        "\n",
        "Get your ngrok authentication token by signing up [here](https://dashboard.ngrok.com/signup). After signing up, copy your auth token [here](https://dashboard.ngrok.com/get-started/your-authtoken) and update the NGROK_AUTH_TOKEN variable with your token."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nEyzcieyHhav"
      },
      "outputs": [],
      "source": [
        "# Uncomment the line below and paste your ngrok auth token\n",
        "NGROK_AUTH_TOKEN = '2M8j0MFxiA90NwVjwSxDIb3vZif_2SkUFXh69AYsab1DN5DVv'\n",
        "TRANSLATE_FILENAME = 'translate.wav'\n",
        "TRANSCRIBE_FILENAME = 'transcribe.wav'\n",
        "WHISPER_MODEL = 'small'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zk57yzgEldFv"
      },
      "source": [
        "## Load WhisperAI model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j28fz90wP5vD"
      },
      "outputs": [],
      "source": [
        "model = whisper.load_model(WHISPER_MODEL)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LAYYG1ulljUl"
      },
      "source": [
        "## Run WhisperAI and host it using ngrok\n",
        "\n",
        "After running, copy the public url provided by ngrok and update WHISPER_BASE_URL in your `.env` file. Below is an example output that you will see. \n",
        "\n",
        "```\n",
        "INFO:werkzeug: * Running on http://127.0.0.1:5000/ (Press CTRL+C to quit)\n",
        " * Running on http://9d55-35-192-10-95.ngrok.io\n",
        " * Traffic stats available on http://127.0.0.1:4040\n",
        "```\n",
        "\n",
        "DO NOT use this url, use the URL provided by the actual output from running the cell below. In this example, you will update your WHISPER_BASE_URL variable with:\n",
        "\n",
        "```\n",
        "WHISPER_BASE_URL=http://9d55-35-192-10-95.ngrok.io\n",
        "```\n",
        "\n",
        "This url will change every time you rerun this cell, so remember to update your `.env` file when that happens."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hCrK1VjzOpJk"
      },
      "outputs": [],
      "source": [
        "app = Flask(__name__)\n",
        "run_with_ngrok(app, auth_token=NGROK_AUTH_TOKEN)\n",
        "\n",
        "\n",
        "@app.route('/', methods=['GET'])\n",
        "def test():\n",
        "    response = {'status':'OK','message':'Test'}\n",
        "    return json.dumps(response)\n",
        "\n",
        "\n",
        "@app.route('/asr', methods=['POST'])\n",
        "def asr():\n",
        "    task = request.args.get('task')\n",
        "    language = request.args.get('language')\n",
        "    audio_data = request.files['audio_file']\n",
        "\n",
        "    if task == 'transcribe':\n",
        "      audio_data.save(TRANSCRIBE_FILENAME)\n",
        "      result = model.transcribe(TRANSCRIBE_FILENAME)\n",
        "      return json.dumps(result)\n",
        "\n",
        "    elif task == 'translate':\n",
        "      audio_data.save(TRANSLATE_FILENAME)\n",
        "      result = model.transcribe(TRANSLATE_FILENAME, language='ja', task='translate')\n",
        "      return json.dumps(result)\n",
        "    \n",
        "    else:\n",
        "      return 'Record not found', 400\n",
        "\n",
        "\n",
        "if __name__ == '__main__':\n",
        "   app.run()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WyTcHX0TmcAt"
      },
      "source": [
        "## Leave the above cell running and the tab open\n",
        "\n",
        "This is to ensure the runtime does not disconnect and shut down the server. \n",
        "\n",
        "When you're done using WhisperAI remember to disconnect the runtime."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}