{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f801f28",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fafb4be8",
   "metadata": {},
   "source": [
    "## Load groundtruths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "917e079a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "import uuid\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from suno_utils.audio import Audio\n",
    "from suno_utils.web.harvest import get_filename\n",
    "\n",
    "BASE_DIR = \"/home/georg/notebooks/tasks/s2t\"\n",
    "DATA_DIR = os.path.join(BASE_DIR, \"sample_data\")\n",
    "AUDIO_DIR = os.path.join(DATA_DIR, \"audio_16k\")\n",
    "\n",
    "AUDIO_DATA = [\n",
    "    {\n",
    "        \"id\": \"a85b9df3-46a4-4461-88ab-c4a52b9fdd96\",\n",
    "        \"filepath\": os.path.join(AUDIO_DIR, \"russia.wav\"),\n",
    "    },\n",
    "    {\n",
    "        \"id\": \"fdeb8ce0-2bca-4b57-ac8b-70c1c8f7202d\",\n",
    "        \"filepath\": os.path.join(AUDIO_DIR, \"georgia.wav\"),\n",
    "    },\n",
    "    {\n",
    "        \"id\": \"a67d8ccc-9bdd-4b1d-be29-88feb1373a4b\",\n",
    "        \"filepath\": os.path.join(AUDIO_DIR, \"spgi_2spk.wav\"),\n",
    "    },\n",
    "    {\n",
    "        \"id\": \"4ac0be6e-d159-4f17-afe3-b1be4e56ef50\",\n",
    "        \"filepath\": os.path.join(AUDIO_DIR, \"spgi_hard_1.wav\"),\n",
    "    },\n",
    "    {\n",
    "        \"id\": \"1bfa8675-ec09-4cf1-a8b6-338c5a6f997d\",\n",
    "        \"filepath\": os.path.join(AUDIO_DIR, \"spgi_hard_2.wav\"),\n",
    "    },\n",
    "    {\n",
    "        \"id\": \"cec77f58-192b-4e56-9fd1-ad5454a23122\",\n",
    "        \"filepath\": os.path.join(AUDIO_DIR, \"spgi_standup.wav\"),\n",
    "    },\n",
    "]\n",
    "\n",
    "with open(os.path.join(DATA_DIR, \"transcripts.json\")) as f:\n",
    "    TRANSCRIPT_CONTAINER = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02fb3c5c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a23e02b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "6e2ded48",
   "metadata": {},
   "source": [
    "## Nemo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "80511a31",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.tasks.asr import transcribe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "e7987020",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "39554345653e4b6ab5a5385fa5e4c8bd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "out = transcribe(\n",
    "    [m[\"filepath\"] for m in AUDIO_DATA], \n",
    "    model_name=\"stt_en_conformer_ctc_large\",\n",
    "#     model_name=\"stt_en_contextnet_1024\",\n",
    "#     model_name=\"stt_en_conformer_transducer_large\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "d907ec6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "for m, text in zip(AUDIO_DATA, out):\n",
    "    TRANSCRIPT_CONTAINER[m[\"id\"]][\"nemo\"] = text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "304cb42a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99e860fa",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1f91b8e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "f8be10dc",
   "metadata": {},
   "source": [
    "## Riva"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72272d6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.web.riva import transcribe_riva"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "a92eca84",
   "metadata": {},
   "outputs": [],
   "source": [
    "out = transcribe_riva([m[\"filepath\"] for m in AUDIO_DATA])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "9e712376",
   "metadata": {},
   "outputs": [],
   "source": [
    "for m, text in zip(AUDIO_DATA, out):\n",
    "    TRANSCRIPT_CONTAINER[m[\"id\"]][\"riva\"] = text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e09c039",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "005846ad",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "1b026f61",
   "metadata": {},
   "source": [
    "## Riva local"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c43047a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !ngc registry model download-version \"nvidia/tao/speechtotext_en_us_conformer:trainable_v4.0\"\n",
    "# !tar -xf speechtotext_en_us_conformer.tlt\n",
    "# !pip install nvidia-eff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a93251ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import os\n",
    "# import yaml\n",
    "# import eff\n",
    "\n",
    "# passphrase = \"tlt_encode\"\n",
    "# model_dir = \"/home/georg/models/speechtotext_en_us_conformer_vtrainable_v4.0\"\n",
    "# with open(os.path.join(model_dir, \"manifest.yaml\")) as f:\n",
    "#     d = yaml.safe_load(f)\n",
    "# encryption = d[\"metadata\"][\"encryption\"]\n",
    "# from_fp = os.path.join(model_dir, \"encrypted_model_weights.ckpt\")\n",
    "# to_fp = os.path.join(model_dir, \"model_weights.ckpt\")\n",
    "# eff.decrypt_file(from_fp, to_fp, passphrase, encryption)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "8715c151",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Exception in thread QueueFeederThread:\n",
      "Traceback (most recent call last):\n",
      "  File \"/usr/lib/python3.8/multiprocessing/queues.py\", line 235, in _feed\n",
      "Exception ignored in:     close()\n",
      "  File \"/usr/lib/python3.8/multiprocessing/connection.py\", line 177, in close\n",
      "<function _ConnectionBase.__del__ at 0x7f1ab8ac4700>\n",
      "Traceback (most recent call last):\n",
      "      File \"/usr/lib/python3.8/multiprocessing/connection.py\", line 132, in __del__\n",
      "self._close()\n",
      "  File \"/usr/lib/python3.8/multiprocessing/connection.py\", line 361, in _close\n",
      "    self._close()\n",
      "  File \"/usr/lib/python3.8/multiprocessing/connection.py\", line 361, in _close\n",
      "    _close(self._handle)\n",
      "OSError: [Errno 9] Bad file descriptor\n",
      "\n",
      "During handling of the above exception, another exception occurred:\n",
      "\n",
      "Traceback (most recent call last):\n",
      "  File \"/usr/lib/python3.8/threading.py\", line 932, in _bootstrap_inner\n",
      "    self.run()\n",
      "  File \"/usr/lib/python3.8/threading.py\", line 870, in run\n",
      "        self._target(*self._args, **self._kwargs)\n",
      "  File \"/usr/lib/python3.8/multiprocessing/queues.py\", line 266, in _feed\n",
      "_close(self._handle)\n",
      "OSError: [Errno 9] Bad file descriptor\n",
      "    queue_sem.release()\n",
      "ValueError: semaphore or lock released too many times\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2022-09-13 13:58:56 mixins:170] Tokenizer SentencePieceTokenizer initialized with 256 tokens\n",
      "[NeMo I 2022-09-13 13:58:56 collections:194] Dataset loaded with 0 files totalling 0.00 hours\n",
      "[NeMo I 2022-09-13 13:58:56 collections:195] 0 files were filtered totalling 0.00 hours\n",
      "[NeMo I 2022-09-13 13:58:56 features:225] PADDING: 0\n",
      "[NeMo I 2022-09-13 13:58:57 audio_preprocessing:491] Numba CUDA SpecAugment kernel is being used\n"
     ]
    }
   ],
   "source": [
    "import yaml\n",
    "import torch\n",
    "from omegaconf.dictconfig import DictConfig\n",
    "from nemo.collections.asr.models.ctc_bpe_models import EncDecCTCModelBPE\n",
    "\n",
    "model_dir = \"/home/georg/models/speechtotext_en_us_conformer_vtrainable_v4.0\"\n",
    "weights_fp = os.path.join(model_dir, \"model_weights.ckpt\")\n",
    "state_dict = torch.load(weights_fp, map_location=\"cpu\")\n",
    "with open(os.path.join(model_dir, \"model_config.yaml\")) as f:\n",
    "    d = yaml.safe_load(f)\n",
    "d[\"tokenizer\"] = {\n",
    "    \"dir\": model_dir, \n",
    "    \"type\": \"bpe\",\n",
    "    \"model_path\": os.path.join(model_dir, d[\"tokenizer\"][\"model_path\"]),\n",
    "    \"vocab_path\": os.path.join(model_dir, d[\"tokenizer\"][\"vocab_path\"]),\n",
    "    \"spe_tokenizer_vocab\": os.path.join(model_dir, d[\"tokenizer\"][\"spe_tokenizer_vocab\"]),\n",
    "}\n",
    "del d[\"train_ds\"]\n",
    "d[\"validation_ds\"][\"manifest_filepath\"] = \"/home/georg/dummy.json\"\n",
    "del d[\"test_ds\"]\n",
    "riva_model = EncDecCTCModelBPE(cfg=DictConfig(d))\n",
    "riva_model.load_state_dict(state_dict);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "bdcbce1d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "50e5d170a24b4684a1c17a5ca0f89f3f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "out = riva_model.transcribe([m[\"filepath\"] for m in AUDIO_DATA])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "c22ee47e",
   "metadata": {},
   "outputs": [],
   "source": [
    "for m, text in zip(AUDIO_DATA, out):\n",
    "    TRANSCRIPT_CONTAINER[m[\"id\"]][\"riva_local\"] = text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc9a9736",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc2e80b5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "bb6a8e81",
   "metadata": {},
   "source": [
    "## Whisper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e49fa23",
   "metadata": {},
   "outputs": [],
   "source": [
    "import whisper\n",
    "\n",
    "# model_sm = whisper.load_model(\"base\")\n",
    "# model = whisper.load_model(\"medium\")\n",
    "# model_lg = whisper.load_model(\"large\")\n",
    "\n",
    "# model_sm = whisper.load_model(\"base.en\")\n",
    "model = whisper.load_model(\"medium.en\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fae8a56",
   "metadata": {},
   "outputs": [],
   "source": [
    "fp = \"sample_data/audio_16k/med_1.wav\"\n",
    "# fp = \"sample_data/audio_16k/med_2.wav\"\n",
    "# fp = \"sample_data/audio_16k/med_3.wav\"\n",
    "\n",
    "audio = whisper.load_audio(fp)\n",
    "audio = whisper.pad_or_trim(audio)\n",
    "mels = whisper.log_mel_spectrogram(audio).to(model.device)\n",
    "options = whisper.DecodingOptions(language=\"en\", without_timestamps=True)\n",
    "result = model.decode(mels, options)\n",
    "print(result.text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39ade352",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60da6a15",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1194fc68",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "6721d98f",
   "metadata": {},
   "source": [
    "## Rev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97ff7ac5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://docs.rev.ai/get-started/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "ac22e215",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"/home/georg/.secrets/secrets.json\") as f:\n",
    "    secrets = json.load(f)\n",
    "REV_AI_ACCESS_TOKEN = secrets[\"REV_AI_ACCESS_TOKEN\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "f9c1a4d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rev_ai import apiclient\n",
    "\n",
    "client = apiclient.RevAiAPIClient(REV_AI_ACCESS_TOKEN)\n",
    "file_job = client.submit_job_local_file(\n",
    "    filename=AUDIO_DATA[0][\"filepath\"], \n",
    "    metadata=\"This_is_some_job_metadata\",\n",
    "    skip_diarization=True,\n",
    "    skip_punctuation=True,\n",
    "    speaker_channels_count=1,\n",
    "    filter_profanity=False,\n",
    "    remove_disfluencies=False,\n",
    "    language=\"en\",\n",
    "    transcriber=\"machine_v2\",\n",
    "    verbatim=True,\n",
    "    rush=None,\n",
    "    test_mode=None,\n",
    "    skip_postprocessing=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "bc8c518f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<JobStatus.TRANSCRIBED: 2>"
      ]
     },
     "execution_count": 89,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "job_details = client.get_job_details(file_job.id)\n",
    "job_details.status"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e48c591b",
   "metadata": {},
   "outputs": [],
   "source": [
    "transcript_json = client.get_transcript_json(file_job.id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "id": "6ab47c66",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "president joe biden backed by the full symbolic power of the western alliance is locked in a showdown with russian president of vladimir putin who is using ukraine as a hostage to try to force the us to renegotiate the settled outcome of the cold war\n"
     ]
    }
   ],
   "source": [
    "out = \" \".join([\" \".join([ee[\"value\"] for ee in e[\"elements\"]]) for e in transcript_json[\"monologues\"]]).lower()\n",
    "print(out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce2445f4",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5bfd46d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4559706",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2037591",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d75f1f7f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c35a733",
   "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
}