{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pPv1dLb5wcmY"
      },
      "source": [
        "# NeMo CTC ONNX ASR — Local Inference Test\n",
        "Quick single-file inference test before pushing to HuggingFace."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ho280TMGwcmd"
      },
      "source": [
        "## 1. Install Dependencies"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D4r4j81-wcmg"
      },
      "outputs": [],
      "source": [
        "#can use onnxruntime and onnxruntime-gpu\n",
        "!pip install -q onnxruntime-gpu soundfile scipy omegaconf pyyaml \"nemo_toolkit[asr]\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "07tQGKt2wcmk"
      },
      "source": [
        "## 2. Mount Drive & Set Paths"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Hdr32L4fwcmp",
        "outputId": "60f6d13e-9cb7-4464-e3f1-baf708a6eed7"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount(\"/content/drive\", force_remount=True).\n"
          ]
        }
      ],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8pDJLHgKwcmq"
      },
      "outputs": [],
      "source": [
        "ONNX_PATH   = \"/content/drive/MyDrive/major_project/onnx_check/model_ctc_quantized.onnx\"\n",
        "CONFIG_PATH = \"/content/drive/MyDrive/major_project/onnx_check/model_config.yaml\"\n",
        "\n",
        "# Single audio file to test — change to any .wav / .flac / .mp3\n",
        "AUDIO_PATH  = \"/content/drive/MyDrive/major_project/onnx_check/nepali_noisy_100%/0125.wav\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7KtwPQVDwcmv"
      },
      "source": [
        "## 3. Load Preprocessor, Vocabulary & ONNX Session"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6bqRRNXZwcmx",
        "outputId": "5c4d9be4-daae-4d83-f37d-f0a53b341de2"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "[NeMo W 2026-03-31 07:39:52 megatron_init:62] Megatron num_microbatches_calculator not found, using Apex version.\n",
            "WARNING:nv_one_logger.api.config:OneLogger: Setting error_handling_strategy to DISABLE_QUIETLY_AND_REPORT_METRIC_ERROR for rank (rank=0) with OneLogger disabled. To override: explicitly set error_handling_strategy parameter.\n",
            "WARNING:nv_one_logger.training_telemetry.api.training_telemetry_provider:No exporters were provided. This means that no telemetry data will be collected.\n",
            "[NeMo W 2026-03-31 07:39:55 nemo_logging:364] /usr/local/lib/python3.12/dist-packages/pydub/utils.py:300: SyntaxWarning: invalid escape sequence '\\('\n",
            "      m = re.match('([su]([0-9]{1,2})p?) \\(([0-9]{1,2}) bit\\)$', token)\n",
            "    \n",
            "[NeMo W 2026-03-31 07:39:55 nemo_logging:364] /usr/local/lib/python3.12/dist-packages/pydub/utils.py:301: SyntaxWarning: invalid escape sequence '\\('\n",
            "      m2 = re.match('([su]([0-9]{1,2})p?)( \\(default\\))?$', token)\n",
            "    \n",
            "[NeMo W 2026-03-31 07:39:55 nemo_logging:364] /usr/local/lib/python3.12/dist-packages/pydub/utils.py:310: SyntaxWarning: invalid escape sequence '\\('\n",
            "      elif re.match('(flt)p?( \\(default\\))?$', token):\n",
            "    \n",
            "[NeMo W 2026-03-31 07:39:55 nemo_logging:364] /usr/local/lib/python3.12/dist-packages/pydub/utils.py:314: SyntaxWarning: invalid escape sequence '\\('\n",
            "      elif re.match('(dbl)p?( \\(default\\))?$', token):\n",
            "    \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Device             : CUDA\n",
            "ORT providers      : ['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider']\n",
            "Preprocessor ready : sample_rate=16000\n",
            "Vocabulary         : 5632 tokens\n",
            "ONNX session ready : ['CUDAExecutionProvider', 'CPUExecutionProvider']\n"
          ]
        }
      ],
      "source": [
        "import numpy as np\n",
        "import onnxruntime as ort\n",
        "import torch\n",
        "import yaml\n",
        "from omegaconf import OmegaConf\n",
        "\n",
        "try:\n",
        "    from nemo.collections.asr.modules import AudioToMelSpectrogramPreprocessor\n",
        "except ModuleNotFoundError:\n",
        "    print(\"The 'nemo_toolkit[asr]' library is not installed. Attempting to install now...\")\n",
        "    !pip install -q nemo_toolkit[asr]\n",
        "    from nemo.collections.asr.modules import AudioToMelSpectrogramPreprocessor # Try importing again\n",
        "\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "print(f\"Device             : {device.upper()}\")\n",
        "print(f\"ORT providers      : {ort.get_available_providers()}\")\n",
        "\n",
        "# Load config\n",
        "try:\n",
        "    conf = OmegaConf.load(CONFIG_PATH)\n",
        "except Exception:\n",
        "    with open(CONFIG_PATH, \"r\", encoding=\"utf-8\") as f:\n",
        "        conf = OmegaConf.create(yaml.safe_load(f))\n",
        "\n",
        "# Preprocessor\n",
        "preprocessor_cfg = OmegaConf.to_container(conf.preprocessor, resolve=True)\n",
        "preprocessor_cfg.pop(\"_target_\", None)\n",
        "preprocessor = AudioToMelSpectrogramPreprocessor(**preprocessor_cfg)\n",
        "preprocessor.eval().to(device)\n",
        "SAMPLE_RATE = preprocessor_cfg[\"sample_rate\"]\n",
        "print(f\"Preprocessor ready : sample_rate={SAMPLE_RATE}\")\n",
        "\n",
        "# Vocabulary\n",
        "vocabulary = (\n",
        "    conf.get(\"aux_ctc\", {}).get(\"decoder\", {}).get(\"vocabulary\", None)\n",
        "    or conf.get(\"decoder\", {}).get(\"vocabulary\", None)\n",
        ")\n",
        "print(f\"Vocabulary         : {len(vocabulary)} tokens\" if vocabulary else \"Vocabulary: NOT FOUND\")\n",
        "\n",
        "# ONNX session\n",
        "providers = [\"CUDAExecutionProvider\", \"CPUExecutionProvider\"] if device == \"cuda\" else [\"CPUExecutionProvider\"]\n",
        "session      = ort.InferenceSession(ONNX_PATH, providers=providers)\n",
        "session_ins  = session.get_inputs()\n",
        "main_input   = next((x for x in session_ins if \"length\" not in x.name.lower()), session_ins[0])\n",
        "length_input = next((x for x in session_ins if \"length\" in x.name.lower()), None)\n",
        "print(f\"ONNX session ready : {session.get_providers()}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pxLr-4zrwcmz"
      },
      "source": [
        "## 4. Run Inference"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DHu736t8wcm3",
        "outputId": "ac8af32f-202a-47bf-f6c6-c216b4231de9"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Audio     : /content/drive/MyDrive/major_project/onnx_check/nepali_noisy_100%/0125.wav\n",
            "Result    : साँचो अदृश्य टिम ल्यर्सन र ल्याथास्टो १९९ेी०९ को उपस्थित पनि भर्चुअल टोलीको अद्वितीय अंश हो<unk>\n"
          ]
        }
      ],
      "source": [
        "import soundfile as sf\n",
        "from scipy.signal import resample_poly\n",
        "\n",
        "def _length_dtype(meta):\n",
        "    return np.int32 if meta and \"int32\" in meta.type else np.int64\n",
        "\n",
        "def decode_ctc(logits, encoded_len, vocab):\n",
        "    greedy   = logits[0].argmax(axis=-1)[: int(encoded_len[0])]\n",
        "    blank_id = logits.shape[-1] - 1\n",
        "    collapsed, prev = [], None\n",
        "    for t in greedy:\n",
        "        t = int(t)\n",
        "        if t == prev or t == blank_id:\n",
        "            prev = t; continue\n",
        "        collapsed.append(t); prev = t\n",
        "    if not vocab:\n",
        "        return str(collapsed)\n",
        "    text = \"\"\n",
        "    for i in collapsed:\n",
        "        if 0 <= i < len(vocab):\n",
        "            tok = vocab[i]\n",
        "            if tok.startswith(\"##\"):   text += tok[2:]\n",
        "            elif tok.startswith(\"▁\"): text += \" \" + tok[1:]\n",
        "            else:                      text += tok\n",
        "    return text.strip().replace(\"▁\", \" \")\n",
        "\n",
        "def transcribe(audio_path: str) -> str:\n",
        "    audio, sr = sf.read(audio_path)\n",
        "    if audio.ndim == 2:       audio = audio.mean(axis=1)\n",
        "    if sr != SAMPLE_RATE:     audio = resample_poly(audio, SAMPLE_RATE, sr)\n",
        "    audio     = np.clip(audio, -1.0, 1.0).astype(np.float32)\n",
        "    audio_len = np.array([audio.shape[0]], dtype=np.int64)\n",
        "\n",
        "    ort_inputs = {}\n",
        "    if len(main_input.shape) == 2:       # raw waveform\n",
        "        ort_inputs[main_input.name] = audio[None, :]\n",
        "        if length_input is not None:\n",
        "            ort_inputs[length_input.name] = audio_len.astype(_length_dtype(length_input))\n",
        "    elif len(main_input.shape) == 3:     # mel-spectrogram\n",
        "        with torch.no_grad():\n",
        "            mel, mel_len = preprocessor(\n",
        "                input_signal=torch.from_numpy(audio[None, :]).to(device),\n",
        "                length=torch.from_numpy(audio_len).to(device),\n",
        "            )\n",
        "        ort_inputs[main_input.name] = mel.cpu().numpy().astype(np.float32)\n",
        "        if length_input is not None:\n",
        "            ort_inputs[length_input.name] = mel_len.cpu().numpy().astype(_length_dtype(length_input))\n",
        "\n",
        "    outputs     = session.run(None, ort_inputs)\n",
        "    logits      = next((x for x in outputs if getattr(x, \"ndim\", 0) == 3), None)\n",
        "    encoded_len = next((x for x in outputs if getattr(x, \"ndim\", 0) == 1), None)\n",
        "    if logits is None:\n",
        "        raise RuntimeError(\"ONNX model returned no 3-D logits tensor.\")\n",
        "    if encoded_len is None:\n",
        "        encoded_len = np.array([logits.shape[1]], dtype=np.int64)\n",
        "\n",
        "    return decode_ctc(logits, encoded_len, vocabulary)\n",
        "\n",
        "# ── Run\n",
        "result = transcribe(AUDIO_PATH)\n",
        "print(f\"\\nAudio     : {AUDIO_PATH}\")\n",
        "print(f\"Result    : {result}\")"
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "MN4duqe5xZp5"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}