indic_asr_inference.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "indic_asr_inference.ipynb",
      "provenance": [],
      "toc_visible": true,
      "authorship_tag": "ABX9TyNAt5QDdurV4AN0lAjt6e3s",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/deterministic-algorithms-lab/indic-asr/blob/main/indic_asr_inference.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DrIVUvArPfqI"
      },
      "source": [
        "!git clone https://github.com/deterministic-algorithms-lab/indic-asr"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "80xUj7Oe4wvH"
      },
      "source": [
        "%cd indic-asr/"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "H4rxNhsEpr-c"
      },
      "source": [
        "!pip install -r requirements.txt\n",
        "!pip install ffmpeg-python"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ztJ9mhrn6jrP"
      },
      "source": [
        "!sudo apt-get install libenchant1c2a\n",
        "!pip install pyenchant"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CfVHLYal4SNQ"
      },
      "source": [
        "## Record Audio to Perform ASR on"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OIMPC3xuQMAO",
        "cellView": "form"
      },
      "source": [
        "#@title Getting Audio from Chrome\n",
        "\"\"\"\n",
        "To write this piece of code I took inspiration/code from a lot of places.\n",
        "It was late night, so I'm not sure how much I created or just copied o.O\n",
        "Here are some of the possible references:\n",
        "https://blog.addpipe.com/recording-audio-in-the-browser-using-pure-html5-and-minimal-javascript/\n",
        "https://stackoverflow.com/a/18650249\n",
        "https://hacks.mozilla.org/2014/06/easy-audio-capture-with-the-mediarecorder-api/\n",
        "https://air.ghost.io/recording-to-an-audio-file-using-html5-and-js/\n",
        "https://stackoverflow.com/a/49019356\n",
        "\"\"\"\n",
        "from IPython.display import HTML, Audio\n",
        "from google.colab.output import eval_js\n",
        "from base64 import b64decode\n",
        "import numpy as np\n",
        "from scipy.io.wavfile import read as wav_read\n",
        "import librosa\n",
        "import io\n",
        "import ffmpeg\n",
        "\n",
        "AUDIO_HTML = \"\"\"\n",
        "<script>\n",
        "var my_div = document.createElement(\"DIV\");\n",
        "var my_p = document.createElement(\"P\");\n",
        "var my_btn = document.createElement(\"BUTTON\");\n",
        "var t = document.createTextNode(\"Press to start recording\");\n",
        "\n",
        "my_btn.appendChild(t);\n",
        "//my_p.appendChild(my_btn);\n",
        "my_div.appendChild(my_btn);\n",
        "document.body.appendChild(my_div);\n",
        "\n",
        "var base64data = 0;\n",
        "var reader;\n",
        "var recorder, gumStream;\n",
        "var recordButton = my_btn;\n",
        "\n",
        "var handleSuccess = function(stream) {\n",
        "  gumStream = stream;\n",
        "  var options = {\n",
        "    //bitsPerSecond: 8000, //chrome seems to ignore, always 48k\n",
        "    mimeType : 'audio/webm;codecs=opus'\n",
        "    //mimeType : 'audio/webm;codecs=pcm'\n",
        "  };            \n",
        "  //recorder = new MediaRecorder(stream, options);\n",
        "  recorder = new MediaRecorder(stream);\n",
        "  recorder.ondataavailable = function(e) {            \n",
        "    var url = URL.createObjectURL(e.data);\n",
        "    var preview = document.createElement('audio');\n",
        "    preview.controls = true;\n",
        "    preview.src = url;\n",
        "    document.body.appendChild(preview);\n",
        "\n",
        "    reader = new FileReader();\n",
        "    reader.readAsDataURL(e.data); \n",
        "    reader.onloadend = function() {\n",
        "      base64data = reader.result;\n",
        "      //console.log(\"Inside FileReader:\" + base64data);\n",
        "    }\n",
        "  };\n",
        "  recorder.start();\n",
        "  };\n",
        "\n",
        "recordButton.innerText = \"Recording... press to stop\";\n",
        "\n",
        "navigator.mediaDevices.getUserMedia({audio: true}).then(handleSuccess);\n",
        "\n",
        "\n",
        "function toggleRecording() {\n",
        "  if (recorder && recorder.state == \"recording\") {\n",
        "      recorder.stop();\n",
        "      gumStream.getAudioTracks()[0].stop();\n",
        "      recordButton.innerText = \"Saving the recording... pls wait!\"\n",
        "  }\n",
        "}\n",
        "\n",
        "// https://stackoverflow.com/a/951057\n",
        "function sleep(ms) {\n",
        "  return new Promise(resolve => setTimeout(resolve, ms));\n",
        "}\n",
        "\n",
        "var data = new Promise(resolve=>{\n",
        "//recordButton.addEventListener(\"click\", toggleRecording);\n",
        "recordButton.onclick = ()=>{\n",
        "toggleRecording()\n",
        "\n",
        "sleep(2000).then(() => {\n",
        "  // wait 2000ms for the data to be available...\n",
        "  // ideally this should use something like await...\n",
        "  //console.log(\"Inside data:\" + base64data)\n",
        "  resolve(base64data.toString())\n",
        "\n",
        "});\n",
        "\n",
        "}\n",
        "});\n",
        "      \n",
        "</script>\n",
        "\"\"\"\n",
        "\n",
        "def get_audio():\n",
        "  display(HTML(AUDIO_HTML))\n",
        "  data = eval_js(\"data\")\n",
        "  binary = b64decode(data.split(',')[1])\n",
        "  \n",
        "  process = (ffmpeg\n",
        "    .input('pipe:0')\n",
        "    .output('pipe:1', format='wav')\n",
        "    .run_async(pipe_stdin=True, pipe_stdout=True, pipe_stderr=True, quiet=True, overwrite_output=True)\n",
        "  )\n",
        "  output, err = process.communicate(input=binary)\n",
        "  \n",
        "  riff_chunk_size = len(output) - 8\n",
        "  # Break up the chunk size into four bytes, held in b.\n",
        "  q = riff_chunk_size\n",
        "  b = []\n",
        "  for i in range(4):\n",
        "      q, r = divmod(q, 256)\n",
        "      b.append(r)\n",
        "\n",
        "  # Replace bytes 4:8 in proc.stdout with the actual size of the RIFF chunk.\n",
        "  riff = output[:4] + bytes(b) + output[8:]\n",
        "\n",
        "  audio, sr = librosa.load(io.BytesIO(riff))\n",
        "\n",
        "  return audio, sr"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KpgBQHO_rGXR",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 96
        },
        "outputId": "9f6d6bd3-3fd3-4f1c-c9dd-788c6ca8513b"
      },
      "source": [
        "audio, sr = get_audio()\n",
        "audio = librosa.resample(audio, sr, 16000)"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "<script>\n",
              "var my_div = document.createElement(\"DIV\");\n",
              "var my_p = document.createElement(\"P\");\n",
              "var my_btn = document.createElement(\"BUTTON\");\n",
              "var t = document.createTextNode(\"Press to start recording\");\n",
              "\n",
              "my_btn.appendChild(t);\n",
              "//my_p.appendChild(my_btn);\n",
              "my_div.appendChild(my_btn);\n",
              "document.body.appendChild(my_div);\n",
              "\n",
              "var base64data = 0;\n",
              "var reader;\n",
              "var recorder, gumStream;\n",
              "var recordButton = my_btn;\n",
              "\n",
              "var handleSuccess = function(stream) {\n",
              "  gumStream = stream;\n",
              "  var options = {\n",
              "    //bitsPerSecond: 8000, //chrome seems to ignore, always 48k\n",
              "    mimeType : 'audio/webm;codecs=opus'\n",
              "    //mimeType : 'audio/webm;codecs=pcm'\n",
              "  };            \n",
              "  //recorder = new MediaRecorder(stream, options);\n",
              "  recorder = new MediaRecorder(stream);\n",
              "  recorder.ondataavailable = function(e) {            \n",
              "    var url = URL.createObjectURL(e.data);\n",
              "    var preview = document.createElement('audio');\n",
              "    preview.controls = true;\n",
              "    preview.src = url;\n",
              "    document.body.appendChild(preview);\n",
              "\n",
              "    reader = new FileReader();\n",
              "    reader.readAsDataURL(e.data); \n",
              "    reader.onloadend = function() {\n",
              "      base64data = reader.result;\n",
              "      //console.log(\"Inside FileReader:\" + base64data);\n",
              "    }\n",
              "  };\n",
              "  recorder.start();\n",
              "  };\n",
              "\n",
              "recordButton.innerText = \"Recording... press to stop\";\n",
              "\n",
              "navigator.mediaDevices.getUserMedia({audio: true}).then(handleSuccess);\n",
              "\n",
              "\n",
              "function toggleRecording() {\n",
              "  if (recorder && recorder.state == \"recording\") {\n",
              "      recorder.stop();\n",
              "      gumStream.getAudioTracks()[0].stop();\n",
              "      recordButton.innerText = \"Saving the recording... pls wait!\"\n",
              "  }\n",
              "}\n",
              "\n",
              "// https://stackoverflow.com/a/951057\n",
              "function sleep(ms) {\n",
              "  return new Promise(resolve => setTimeout(resolve, ms));\n",
              "}\n",
              "\n",
              "var data = new Promise(resolve=>{\n",
              "//recordButton.addEventListener(\"click\", toggleRecording);\n",
              "recordButton.onclick = ()=>{\n",
              "toggleRecording()\n",
              "\n",
              "sleep(2000).then(() => {\n",
              "  // wait 2000ms for the data to be available...\n",
              "  // ideally this should use something like await...\n",
              "  //console.log(\"Inside data:\" + base64data)\n",
              "  resolve(base64data.toString())\n",
              "\n",
              "});\n",
              "\n",
              "}\n",
              "});\n",
              "      \n",
              "</script>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O0qRlsnE3P27"
      },
      "source": [
        "## Download and save weights"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VDyyd740De1U"
      },
      "source": [
        "!rm -r facebook/wav2vec2-base-960h_14/"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8UhGHW8F3OJv"
      },
      "source": [
        "import wandb\n",
        "api = wandb.Api()\n",
        "run = api.run(\"interspeech-asr/wav2vec2.0/3bh42rgp/\")\n",
        "run.file(\"facebook/wav2vec2-base-960h_14/pytorch_model.bin\").download()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Sa_3HtWO4imt"
      },
      "source": [
        "## Run Model & get Predictions"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Eyme2P1Y6eK4"
      },
      "source": [
        "%cd src"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BbJNH1F9P8Al"
      },
      "source": [
        "from configs import config, get_all_params_dict\n",
        "from model import get_model\n",
        "from tokenizer import Wav2Vec2Tok"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ohXeGYMM82HK"
      },
      "source": [
        "config.prev_checkpoint = \"../facebook/wav2vec2-base-960h_14/\""
      ],
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gh8QO-rW_f-r"
      },
      "source": [
        "import torch"
      ],
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PDy77qsAQCfq"
      },
      "source": [
        "all_params_dict = get_all_params_dict(config)\n",
        "tokenizer = Wav2Vec2Tok.from_pretrained(config.model)\n",
        "model = get_model(tokenizer, n_langs=1)\n",
        "model.load_state_dict(torch.load(config.prev_checkpoint+\"pytorch_model.bin\", map_location=torch.device(\"cpu\")), strict=False)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6I2rsXMPUFhw"
      },
      "source": [
        "speech = tokenizer([audio], return_tensors='pt', padding='longest').input_values\n",
        "logits = model(speech).logits\n",
        "predicted_ids = torch.argmax(logits, dim=-1)\n",
        "predicted_text = tokenizer.batch_decode(predicted_ids)\n",
        "prediction = tokenizer.revert_transliteration(predicted_text)"
      ],
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "N_5b4tpZv1g6"
      },
      "source": [
        "print(prediction)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "muDgyKZyE8zu"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}