{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "3459a8eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import grpc\n",
    "import time\n",
    "import riva_api.riva_audio_pb2 as ra\n",
    "import riva_api.riva_asr_pb2 as rasr\n",
    "import riva_api.riva_asr_pb2_grpc as rasr_srv\n",
    "import wave\n",
    "import threading\n",
    "\n",
    "\n",
    "def _get_single_transcripts(filepath):\n",
    "    channel = grpc.insecure_channel(\"localhost:50051\")\n",
    "    client = rasr_srv.RivaSpeechRecognitionStub(channel)\n",
    "    config = rasr.RecognitionConfig(\n",
    "        encoding=ra.AudioEncoding.LINEAR_PCM,\n",
    "        sample_rate_hertz=44_100,\n",
    "        language_code=\"en-US\",\n",
    "        max_alternatives=1,\n",
    "        enable_automatic_punctuation=False,\n",
    "        audio_channel_count=1\n",
    "    )\n",
    "    with open(filepath, 'rb') as fh:\n",
    "        data = fh.read()\n",
    "    request = rasr.RecognizeRequest(config=config, audio=data)\n",
    "    response = client.Recognize(request)\n",
    "    alternatives = response.results[0].alternatives\n",
    "    out = alternatives[0].transcript if len(alternatives) > 0 else \"\"\n",
    "    return out\n",
    "\n",
    "\n",
    "def _get_batch_transcripts(indexed_filepaths, result):\n",
    "    for idx, filepath in indexed_filepaths:\n",
    "        result[idx] = _get_single_transcripts(filepath)\n",
    "        \n",
    "def transcribe_riva(filepaths):\n",
    "    input_is_str = False\n",
    "    if isinstance(filepaths, str):\n",
    "        input_is_str = True\n",
    "        filepaths = [filepaths]\n",
    "    n_threads = len(filepaths) if len(filepaths) <= 20 else 20\n",
    "    threads = [None] * n_threads\n",
    "    filepath_chunks = [[] for _ in range(n_threads)]\n",
    "    results = [None] * len(filepaths)\n",
    "    for n, filepath in enumerate(filepaths):\n",
    "        filepath_chunks[n % n_threads].append((n, filepath))\n",
    "    for i in range(len(threads)):\n",
    "        threads[i] = threading.Thread(target=_get_batch_transcripts, args=(filepath_chunks[i], results))\n",
    "        threads[i].start()\n",
    "    for i in range(len(threads)):\n",
    "        threads[i].join()\n",
    "    if input_is_str:\n",
    "        results = results[0]\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "ebebe808",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "DATA_DIR = \"/mnt/data-ssd-1/data/gdata-mate/pipeline/\"\n",
    "\n",
    "with open(DATA_DIR + \"artifacts/04_segment_meta_distribution.jsonl\") as f:\n",
    "    d = [json.loads(l) for l in f.read().strip().split(\"\\n\")]\n",
    "    \n",
    "filepaths = [DATA_DIR + e[\"uri\"] for e in d]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "03b08900",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"Without finishing it off, so I'm going to wrap it up with the last few days of the twelve days of Christmas countdown right now, In one foul swot, one fell swoop, one fowl scoop. Whatever sweep it is, I probably lean towards fowl swoop, being as a bird swoops, that may be a f swoop. I much sh, anyway, so bang, the two go Moses, twelve days of Christmas countdown. \""
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "transcribe_riva(filepaths[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ce95b51e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"without finishing it off so i'm going to wrap it up with the last few days of the twelve days of christmas countdown right now in one foul swot one fell swoop one fowl scoop whatever sweep it is i probably lean towards fowl swoop being as a bird swoops that may be a f swoop i much sh anyway so bang the two go moses twelve days of christmas countdown\""
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19e426f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# do on all"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff5eadcb",
   "metadata": {},
   "outputs": [],
   "source": [
    "t0 = time.time()\n",
    "hyps = transcribe_riva(filepaths)\n",
    "t1 = time.time()\n",
    "print(round((t1 - t0) / 60, 1), \"minutes runtime\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "dc5af14e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.14325885133653354"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from suno_utils.utils.metrics import get_wer_bulk\n",
    "refs = [e[\"transcript\"][\"text_normalized\"] for e in d]\n",
    "get_wer_bulk(refs, hyps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11cdf5b0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfc1016d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "c4a03195",
   "metadata": {},
   "source": [
    "## run on hard subset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ee87bc65",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"nemo_suno_aus_samples/transcripts.json\") as f:\n",
    "    hard_examples = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "0757652c",
   "metadata": {},
   "outputs": [],
   "source": [
    "refs = []\n",
    "fps = []\n",
    "for k, v in hard_examples.items():\n",
    "    refs.append(v[\"text_norm\"])\n",
    "    fps.append(f\"nemo_suno_aus_samples/{k}.wav\")\n",
    "hyps = transcribe_riva(fps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "90518d78",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.3157051282051282"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from suno_utils.utils.metrics import get_wer_bulk\n",
    "get_wer_bulk(refs, hyps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f48049bb",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "daec74a4",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d3ac382",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "fe02d42f",
   "metadata": {},
   "source": [
    "## Playground"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1381bd28",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"and it's just about keeping peaceful so i like that i wonder you created what are those three things go by let's run through them again relax look around make a call do it do it and the other part that i think that this is again i'm just going to keep sort betting a lot about businesses why i love jock least extreme ownership principles is if you could always pitch yourself in a war zone starts to give you a lot of clarity that maybe what you're going through at the moment within your\""
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "transcribe_riva(\"nemo_suno_aus_samples/13fc3f02-f864-4551-91a4-480109b4ae62.wav\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca4ba44c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "526d5c29",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aae8ac31",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69ea85d7",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "329729c6",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}