diff --git a/.github/workflows/export-nvidia-asr-model.yml b/.github/workflows/export-nvidia-asr-model.yml
index dc098d1..8ab9c4b 100644
--- a/.github/workflows/export-nvidia-asr-model.yml
+++ b/.github/workflows/export-nvidia-asr-model.yml
@@ -65,8 +65,12 @@ jobs:
elif [[ -z "$BASE_SHA" ]] || ! git cat-file -e "${BASE_SHA}^{commit}"; then
force=true
elif git diff --quiet "$BASE_SHA" "$SOURCE_REF" -- \
- 'asr/rnnt/python/*.py' \
- 'asr/rnnt/python/**/*.py' \
+ 'asr/parakeet/model_export/**' \
+ 'asr/nemotron/model_export/**' \
+ 'asr/parakeet/requirements.txt' \
+ 'asr/nemotron/requirements.txt' \
+ 'common/model_export/nemo_export.py' \
+ 'common/model_export/requirements-nemo.txt' \
'.github/workflows/export-nvidia-asr-model.yml' \
'.github/workflows/nvidia-asr.yml'; then
force=false
@@ -194,11 +198,7 @@ jobs:
set -euxo pipefail
python -m pip install --upgrade pip
python -m pip install --index-url https://download.pytorch.org/whl/cpu torch
- python -m pip install \
- onnx \
- onnxscript \
- soundfile \
- "nemo_toolkit[asr] @ git+https://github.com/NVIDIA/NeMo.git@main"
+ python -m pip install -r common/model_export/requirements-nemo.txt
- name: Export NVIDIA ASR ONNX
if: steps.selection.outputs.export == 'true'
@@ -211,7 +211,6 @@ jobs:
MODEL_SLUG: ${{ inputs.model_slug }}
MODEL_TYPE: ${{ inputs.model_type }}
PRECISION: ${{ inputs.precision }}
- PYTHONPATH: asr/rnnt/python
SOURCE_REF: ${{ inputs.source_ref }}
run: |
set -uo pipefail
@@ -227,7 +226,7 @@ jobs:
dtype="$(export_dtype "$1")"
case "$MODEL_TYPE" in
parakeet)
- python asr/rnnt/python/tools/export_parakeet_tdt_onnx.py \
+ python asr/parakeet/model_export/export_parakeet.py \
--model "$MODEL_ID" \
--output "$MODEL_OUTPUT" \
--device cpu \
@@ -235,7 +234,7 @@ jobs:
--external-data
;;
nemotron)
- python asr/rnnt/python/tools/export_nemotron_onnx.py \
+ python asr/nemotron/model_export/export_nemotron.py \
--model "$MODEL_ID" \
--output "$MODEL_OUTPUT" \
--device cpu \
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 4ac94bd..26d96bf 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -99,7 +99,10 @@ add_subdirectory(tests/base)
add_subdirectory(tests/tokenizer)
add_subdirectory(image_gen/flux2)
add_subdirectory(image_gen/pixeldit)
-add_subdirectory(asr/rnnt)
+add_subdirectory(asr/common)
+add_subdirectory(tests/asr_cli)
+add_subdirectory(asr/parakeet)
+add_subdirectory(asr/nemotron)
add_subdirectory(vision/sam2)
add_subdirectory(asr/whisper)
add_subdirectory(asr/qwen3)
diff --git a/README.md b/README.md
index 5f111ff..133318f 100644
--- a/README.md
+++ b/README.md
@@ -17,8 +17,8 @@ DIN Deploy is a collection of practical samples for exporting and running local
| Model | Hugging Face ID | Docs |
| --- | --- | --- |
| OpenAI Whisper | `openai/whisper-tiny`
`openai/whisper-base`
`openai/whisper-small`
`openai/whisper-medium`
`openai/whisper-large-v3`
`openai/whisper-large-v3-turbo` | [Whisper](asr/whisper/README.md) |
-| NVIDIA Parakeet TDT 0.6B v3 | `nvidia/parakeet-tdt-0.6b-v3` | [RNNT](asr/rnnt/README.md) |
-| NVIDIA Nemotron 3.5 ASR Streaming 0.6B | `nvidia/nemotron-3.5-asr-streaming-0.6b` | [RNNT](asr/rnnt/README.md) |
+| NVIDIA Parakeet TDT 0.6B v3 | `nvidia/parakeet-tdt-0.6b-v3` | [Parakeet](asr/parakeet/README.md) |
+| NVIDIA Nemotron 3.5 ASR Streaming 0.6B | `nvidia/nemotron-3.5-asr-streaming-0.6b` | [Nemotron ASR](asr/nemotron/README.md) |
| Qwen3 ASR | `Qwen/Qwen3-ASR-0.6B-hf`
`Qwen/Qwen3-ASR-1.7B-hf` | [Qwen3](asr/qwen3/README.md) |
| Qwen3 Forced Aligner | `Qwen/Qwen3-ForcedAligner-0.6B-hf` | [Qwen3](asr/qwen3/README.md) |
diff --git a/asr/common/CMakeLists.txt b/asr/common/CMakeLists.txt
new file mode 100644
index 0000000..ffada8d
--- /dev/null
+++ b/asr/common/CMakeLists.txt
@@ -0,0 +1,2 @@
+add_library(din_asr_common INTERFACE)
+target_include_directories(din_asr_common INTERFACE "${CMAKE_CURRENT_SOURCE_DIR}")
diff --git a/asr/common/asr_common.h b/asr/common/asr_common.h
new file mode 100644
index 0000000..74bd695
--- /dev/null
+++ b/asr/common/asr_common.h
@@ -0,0 +1,78 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+
+namespace din::asr::common
+{
+
+// --- file helpers --------------------------------------------------------
+
+inline std::string ReadText(const std::string& path)
+{
+ std::ifstream stream(path);
+ if (!stream)
+ {
+ throw std::runtime_error("failed to open: " + path);
+ }
+ std::ostringstream buffer;
+ buffer << stream.rdbuf();
+ return buffer.str();
+}
+
+inline std::string RequirePath(const std::filesystem::path& path, std::string_view name)
+{
+ if (!std::filesystem::exists(path))
+ {
+ throw std::runtime_error("missing " + std::string{name} + ": " + path.string());
+ }
+ return path.string();
+}
+
+inline std::vector ReadFloatFile(const std::filesystem::path& path)
+{
+ std::ifstream stream(path, std::ios::binary);
+ if (!stream)
+ {
+ return {};
+ }
+ stream.seekg(0, std::ios::end);
+ const auto size = stream.tellg();
+ stream.seekg(0, std::ios::beg);
+ std::vector values(static_cast(size) / sizeof(float));
+ stream.read(reinterpret_cast(values.data()), size);
+ return values;
+}
+
+// --- timing helper -------------------------------
+
+inline double SecondsSince(std::chrono::steady_clock::time_point start)
+{
+ return std::chrono::duration(std::chrono::steady_clock::now() - start).count();
+}
+
+// --- Shared command-line front end --------------------------------------
+
+struct TranscriptionOptions
+{
+ std::string timestamps = "none";
+};
+
+} // namespace din::asr::common
diff --git a/asr/rnnt/cpp/asr_common.h b/asr/common/cli.h
similarity index 72%
rename from asr/rnnt/cpp/asr_common.h
rename to asr/common/cli.h
index eb2d6d7..c956fb8 100644
--- a/asr/rnnt/cpp/asr_common.h
+++ b/asr/common/cli.h
@@ -3,82 +3,13 @@
#pragma once
-#include
-#include
-#include
-#include
-#include
-#include
-#include
-#include
-#include
-#include
-#include
-#include
-#include
-#include
-#include
-#include
-
-#include "ort_session.h"
+#include "asr_common.h"
#include
-#include
+#include
namespace din::asr::common
{
-// --- file helpers --------------------------------------------------------
-
-inline std::string ReadText(const std::string& path)
-{
- std::ifstream stream(path);
- if (!stream)
- {
- throw std::runtime_error("failed to open: " + path);
- }
- std::ostringstream buffer;
- buffer << stream.rdbuf();
- return buffer.str();
-}
-
-inline std::string RequirePath(const std::filesystem::path& path, std::string_view name)
-{
- if (!std::filesystem::exists(path))
- {
- throw std::runtime_error("missing " + std::string{name} + ": " + path.string());
- }
- return path.string();
-}
-
-inline std::vector ReadFloatFile(const std::filesystem::path& path)
-{
- std::ifstream stream(path, std::ios::binary);
- if (!stream)
- {
- return {};
- }
- stream.seekg(0, std::ios::end);
- const auto size = stream.tellg();
- stream.seekg(0, std::ios::beg);
- std::vector values(static_cast(size) / sizeof(float));
- stream.read(reinterpret_cast(values.data()), size);
- return values;
-}
-
-// --- timing helper -------------------------------
-
-inline double SecondsSince(std::chrono::steady_clock::time_point start)
-{
- return std::chrono::duration(std::chrono::steady_clock::now() - start).count();
-}
-
-// --- Shared command-line front end --------------------------------------
-
-struct TranscriptionOptions
-{
- std::string timestamps = "none";
-};
-
struct CliSpec
{
std::string fallback_name; // executable name when argv[0] is unavailable
@@ -102,6 +33,7 @@ int RunCli(int argc, char** argv, const CliSpec& spec)
argparse::ArgumentParser parser(spec.fallback_name);
parser.add_description("Transcribe audio with exported " + spec.artifact_label + " ONNX artifacts.");
+ parser.add_argument("--json").default_value(false).implicit_value(true).help("Print transcription as JSON.");
parser.add_argument("audiofile").help("Audio file to transcribe.");
parser.add_argument("--model-dir")
.default_value(config.model_dir.string())
@@ -207,16 +139,25 @@ int RunCli(int argc, char** argv, const CliSpec& spec)
const auto load_seconds = SecondsSince(load_start);
const auto result = pipeline.TranscribeFile(audio_file);
- pipeline.Print(std::cout, result, options);
-
- std::cout << "\nprovider: " << provider << '\n';
- std::cout << "audio: " << result.audio_seconds << "s\n";
- std::cout << "load: " << load_seconds << "s\n";
- std::cout << "transcribe: " << result.transcribe_seconds << "s\n";
- std::cout << "encode: " << result.encode_seconds << "s\n";
- std::cout << "greedy: " << result.greedy_seconds << "s\n";
- std::cout << "speed: " << (result.audio_seconds / result.transcribe_seconds) << "x\n";
- std::cout << "total: " << SecondsSince(total_start) << "s\n";
+ const bool json_output = parser.get("--json");
+ if (json_output)
+ std::cout << nlohmann::json{{"transcription", result.text},
+ {"audio_seconds", result.audio_seconds},
+ {"transcribe_seconds", result.transcribe_seconds}}
+ .dump()
+ << '\n';
+ else
+ pipeline.Print(std::cout, result, options);
+ auto& timing = json_output ? std::cerr : std::cout;
+
+ timing << "\nprovider: " << provider << '\n';
+ timing << "audio: " << result.audio_seconds << "s\n";
+ timing << "load: " << load_seconds << "s\n";
+ timing << "transcribe: " << result.transcribe_seconds << "s\n";
+ timing << "encode: " << result.encode_seconds << "s\n";
+ timing << "greedy: " << result.greedy_seconds << "s\n";
+ timing << "speed: " << (result.audio_seconds / result.transcribe_seconds) << "x\n";
+ timing << "total: " << SecondsSince(total_start) << "s\n";
}
catch (const std::exception& exception)
{
diff --git a/asr/nemotron/CMakeLists.txt b/asr/nemotron/CMakeLists.txt
new file mode 100644
index 0000000..4cd572f
--- /dev/null
+++ b/asr/nemotron/CMakeLists.txt
@@ -0,0 +1,6 @@
+add_din_shared_library(din_asr_nemotron nemotron.cpp)
+target_include_directories(din_asr_nemotron PUBLIC "${CMAKE_CURRENT_SOURCE_DIR}")
+target_link_libraries(din_asr_nemotron PUBLIC din_asr_common din_common_ort din_common_io nlohmann_json::nlohmann_json)
+
+add_din_executable(din_asr_nemotron_cli main.cpp)
+target_link_libraries(din_asr_nemotron_cli PRIVATE din_asr_nemotron argparse)
diff --git a/asr/nemotron/README.md b/asr/nemotron/README.md
new file mode 100644
index 0000000..ba13f28
--- /dev/null
+++ b/asr/nemotron/README.md
@@ -0,0 +1,39 @@
+# Nemotron 3.5 ASR
+
+NeMo → standard ONNX → C++ inference through ONNX Runtime (CPU or TensorRT RTX).
+
+## Export
+
+Python 3.12+ is needed only for export and reference validation.
+
+```bash
+python -m pip install -r asr/nemotron/requirements.txt
+python asr/nemotron/model_export/export_nemotron.py --output artifacts/nemotron/onnx --device cpu --external-data
+```
+
+## Build and run
+
+```bash
+cmake --preset linux-x64
+cmake --build --preset linux-x64-release --target din_asr_nemotron_cli
+out/build/linux-x64/bin/Release/din_asr_nemotron_cli audio.wav --model-dir artifacts/nemotron/onnx --provider cpu --lang-id en-US
+```
+
+Use `--json` for a JSON transcript, or `--timestamps token` for token timestamps.
+The existing C++ target names and exported artifact formats are preserved.
+
+## Validate
+
+Compare one real recording against NeMo; add `--truth transcript.txt` to measure WER
+against a human transcript, and `--max-wer 0.1` to enforce a threshold.
+
+```bash
+python asr/nemotron/model_export/validate_nemotron.py --audio audio.wav \
+ --cli out/build/linux-x64/bin/Release/din_asr_nemotron_cli \
+ --model-dir artifacts/nemotron/onnx --lang-id en-US
+```
+
+The Python inference package, native bindings, dataset evaluation, plotting, and
+quantization tools have been removed. Run inference with the C++ CLI; prepare
+models with the scripts above. Shared NeMo export and validation helpers live in
+`common/model_export/`; shared C++ ASR utilities live in `asr/common/`.
diff --git a/asr/rnnt/cpp/nemotron_main.cpp b/asr/nemotron/main.cpp
similarity index 98%
rename from asr/rnnt/cpp/nemotron_main.cpp
rename to asr/nemotron/main.cpp
index d5d4a32..2b4cea6 100644
--- a/asr/rnnt/cpp/nemotron_main.cpp
+++ b/asr/nemotron/main.cpp
@@ -1,7 +1,7 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
-#include "asr_common.h"
+#include "cli.h"
#include "nemotron.h"
int main(int argc, char** argv)
diff --git a/asr/rnnt/python/tools/export_nemotron_onnx.py b/asr/nemotron/model_export/export_nemotron.py
similarity index 90%
rename from asr/rnnt/python/tools/export_nemotron_onnx.py
rename to asr/nemotron/model_export/export_nemotron.py
index 4031d2e..2a0ac55 100644
--- a/asr/rnnt/python/tools/export_nemotron_onnx.py
+++ b/asr/nemotron/model_export/export_nemotron.py
@@ -12,25 +12,16 @@
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
-from typing import Any, TypeVar
+from typing import TypeVar
import torch
from torch import nn
-PYTHON_DIR = Path(__file__).resolve().parents[1]
-if str(PYTHON_DIR) not in sys.path:
- sys.path.insert(0, str(PYTHON_DIR))
+sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "common" / "model_export"))
+from nemo_export import configure_encoder_sdpa, export_nemo_preprocessor, load_nemo_model
-from rnnt.nemotron_asr.nemo_backend import (
- DEFAULT_ATT_CONTEXT_SIZE,
- DEFAULT_MODEL_ID,
- load_nemo_model,
- prompt_dictionary,
- prompt_id,
- write_tokenizer_artifacts,
-)
-
-from nemo_preprocessor_export import export_nemo_preprocessor
+DEFAULT_MODEL_ID = "nvidia/nemotron-3.5-asr-streaming-0.6b"
+DEFAULT_ATT_CONTEXT_SIZE = (56, 3)
ARTIFACTS_DIR = Path("artifacts") / "nemotron"
ONNX_DIR = ARTIFACTS_DIR / "onnx"
@@ -235,12 +226,10 @@ def main() -> None:
lambda: load_nemo_model(
args.model,
device=device,
- target_lang=args.target_lang,
- att_context_size=args.att_context_size,
dtype=dtype,
- configure_for_export=True,
),
)
+ configure_model(model, args.target_lang, args.att_context_size, export=True)
if hasattr(model.encoder, "setup_streaming_params"):
model.encoder.setup_streaming_params(
chunk_size=args.chunk_size,
@@ -248,7 +237,6 @@ def main() -> None:
shift_size=args.chunk_size,
)
model.eval()
- model, _ = restore_quantized_modelopt_checkpoint(model, args)
graphs = parse_graph_names(args.graphs)
if "preprocessor" in graphs:
@@ -431,19 +419,6 @@ def export_decoder_init(
)
-def restore_quantized_modelopt_checkpoint(model, args: argparse.Namespace) -> tuple[object, dict[str, Any]]:
- if args.quantized_modelopt_path is None:
- return model, {}
- from rnnt.validation.quantization import restore_modelopt_quantized_model
-
- restored_model, info = restore_modelopt_quantized_model(
- model,
- args.quantized_modelopt_path,
- summary_path=args.output / "quant_summary.txt",
- )
- return restored_model, {"modelopt_checkpoint": info}
-
-
def write_metadata(
model: nn.Module,
output: Path,
@@ -563,12 +538,6 @@ def parse_args() -> argparse.Namespace:
)
parser.add_argument("--strip-lang-tags", "--strip_lang_tags", action=argparse.BooleanOptionalAction, default=True)
parser.add_argument("--max-symbols-per-step", type=int, default=10)
- parser.add_argument(
- "--quantized-modelopt-path",
- type=Path,
- default=None,
- help="Restore a saved ModelOpt encoder checkpoint before exporting.",
- )
parser.add_argument(
"--graphs",
default="preprocessor,encoder_step,predict_step",
@@ -577,6 +546,44 @@ def parse_args() -> argparse.Namespace:
return parser.parse_args()
+def prompt_dictionary(nemo_model: nn.Module) -> dict[str, int]:
+ for key in ("train_ds", "validation_ds", "test_ds"):
+ data = nemo_model.cfg.get(key)
+ if data is not None and "prompt_dictionary" in data:
+ return {str(lang): int(prompt_id) for lang, prompt_id in data.prompt_dictionary.items()}
+ raise KeyError("prompt_dictionary")
+
+
+def prompt_id(nemo_model: nn.Module, target_lang: str) -> int:
+ prompts = prompt_dictionary(nemo_model)
+ if target_lang not in prompts:
+ raise ValueError(f"Unsupported target language: {target_lang}")
+ return prompts[target_lang]
+
+
+def write_tokenizer_artifacts(nemo_model: nn.Module, output: Path) -> None:
+ tokenizer = nemo_model.tokenizer.tokenizer
+ (output / "tokenizer.model").write_bytes(tokenizer.serialized_model_proto())
+ pieces = [tokenizer.id_to_piece(index) for index in range(tokenizer.get_piece_size())]
+ (output / "vocab.txt").write_text("\n".join(pieces) + "\n", encoding="utf-8")
+
+
+def configure_model(
+ model: nn.Module, target_lang: str, att_context_size: tuple[int, int], *, export: bool = False
+) -> None:
+ if hasattr(model, "set_inference_prompt"):
+ model.set_inference_prompt(target_lang)
+ if export and hasattr(model, "set_export_config"):
+ model.set_export_config({"cache_support": True})
+ if hasattr(model.encoder, "set_default_att_context_size"):
+ model.encoder.set_default_att_context_size(list(att_context_size))
+ if export:
+ configure_encoder_sdpa(model.encoder)
+ if hasattr(model.encoder, "export_cache_support"):
+ model.encoder.export_cache_support = True
+ model.encoder.eval()
+
+
if __name__ == "__main__":
with torch.inference_mode():
main()
diff --git a/asr/nemotron/model_export/validate_nemotron.py b/asr/nemotron/model_export/validate_nemotron.py
new file mode 100644
index 0000000..cd27e6b
--- /dev/null
+++ b/asr/nemotron/model_export/validate_nemotron.py
@@ -0,0 +1,32 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+import re
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "common" / "model_export"))
+from export_nemotron import DEFAULT_MODEL_ID, configure_model
+from nemo_export import load_nemo_model
+from validate_asr import validate, validation_parser
+
+
+def reference(args):
+ from nemo.collections.asr.models.rnnt_bpe_models_prompt import RNNTPromptTranscribeConfig
+
+ metadata = json.loads((args.model_dir / "metadata.json").read_text(encoding="utf-8"))
+ model = load_nemo_model(args.model, device=args.device)
+ configure_model(model, args.lang_id, tuple(metadata["att_context_size"]))
+ config = RNNTPromptTranscribeConfig(batch_size=1, verbose=False, target_lang=args.lang_id)
+ config.use_lhotse = False
+ result = model.transcribe([str(args.audio.resolve())], override_config=config)[0]
+ text = str(getattr(result, "text", result)).strip()
+ return re.sub(r"<[^>\s]+>", "", text).strip() if metadata.get("strip_lang_tags", True) else text
+
+
+if __name__ == "__main__":
+ parser = validation_parser(DEFAULT_MODEL_ID)
+ parser.add_argument("--lang-id", default="auto", help="Use the same target language for NeMo and C++.")
+ args = parser.parse_args()
+ validate(args, reference, ["--lang-id", args.lang_id])
diff --git a/asr/rnnt/cpp/nemotron.cpp b/asr/nemotron/nemotron.cpp
similarity index 100%
rename from asr/rnnt/cpp/nemotron.cpp
rename to asr/nemotron/nemotron.cpp
diff --git a/asr/rnnt/cpp/nemotron.h b/asr/nemotron/nemotron.h
similarity index 100%
rename from asr/rnnt/cpp/nemotron.h
rename to asr/nemotron/nemotron.h
diff --git a/asr/nemotron/requirements.txt b/asr/nemotron/requirements.txt
new file mode 100644
index 0000000..eb67fa9
--- /dev/null
+++ b/asr/nemotron/requirements.txt
@@ -0,0 +1 @@
+-r ../../common/model_export/requirements-nemo.txt
diff --git a/asr/parakeet/CMakeLists.txt b/asr/parakeet/CMakeLists.txt
new file mode 100644
index 0000000..21ca749
--- /dev/null
+++ b/asr/parakeet/CMakeLists.txt
@@ -0,0 +1,6 @@
+add_din_shared_library(din_asr_parakeet_tdt parakeet_tdt.cpp)
+target_include_directories(din_asr_parakeet_tdt PUBLIC "${CMAKE_CURRENT_SOURCE_DIR}")
+target_link_libraries(din_asr_parakeet_tdt PUBLIC din_asr_common din_common_ort din_common_io nlohmann_json::nlohmann_json)
+
+add_din_executable(din_asr_parakeet_tdt_cli main.cpp)
+target_link_libraries(din_asr_parakeet_tdt_cli PRIVATE din_asr_parakeet_tdt argparse)
diff --git a/asr/parakeet/README.md b/asr/parakeet/README.md
new file mode 100644
index 0000000..9b2cc57
--- /dev/null
+++ b/asr/parakeet/README.md
@@ -0,0 +1,39 @@
+# Parakeet TDT
+
+NeMo → standard ONNX → C++ inference through ONNX Runtime (CPU or TensorRT RTX).
+
+## Export
+
+Python 3.12+ is needed only for export and reference validation.
+
+```bash
+python -m pip install -r asr/parakeet/requirements.txt
+python asr/parakeet/model_export/export_parakeet.py --output artifacts/parakeet/onnx --device cpu --external-data
+```
+
+## Build and run
+
+```bash
+cmake --preset linux-x64
+cmake --build --preset linux-x64-release --target din_asr_parakeet_tdt_cli
+out/build/linux-x64/bin/Release/din_asr_parakeet_tdt_cli audio.wav --model-dir artifacts/parakeet/onnx --provider cpu
+```
+
+Use `--json` for a JSON transcript, or `--timestamps token` for token timestamps.
+The existing C++ target names and exported artifact formats are preserved.
+
+## Validate
+
+Compare one real recording against NeMo; add `--truth transcript.txt` to measure WER
+against a human transcript, and `--max-wer 0.1` to enforce a threshold.
+
+```bash
+python asr/parakeet/model_export/validate_parakeet.py --audio audio.wav \
+ --cli out/build/linux-x64/bin/Release/din_asr_parakeet_tdt_cli \
+ --model-dir artifacts/parakeet/onnx
+```
+
+The Python inference package, native bindings, dataset evaluation, plotting, and
+quantization tools have been removed. Run inference with the C++ CLI; prepare
+models with the scripts above. Shared NeMo export and validation helpers live in
+`common/model_export/`; shared C++ ASR utilities live in `asr/common/`.
diff --git a/asr/rnnt/cpp/parakeet_tdt_main.cpp b/asr/parakeet/main.cpp
similarity index 96%
rename from asr/rnnt/cpp/parakeet_tdt_main.cpp
rename to asr/parakeet/main.cpp
index 8571a2b..7730a94 100644
--- a/asr/rnnt/cpp/parakeet_tdt_main.cpp
+++ b/asr/parakeet/main.cpp
@@ -1,7 +1,7 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
-#include "asr_common.h"
+#include "cli.h"
#include "parakeet_tdt.h"
int main(int argc, char** argv)
diff --git a/asr/rnnt/python/tools/export_parakeet_tdt_onnx.py b/asr/parakeet/model_export/export_parakeet.py
similarity index 70%
rename from asr/rnnt/python/tools/export_parakeet_tdt_onnx.py
rename to asr/parakeet/model_export/export_parakeet.py
index 791af75..9d1519c 100644
--- a/asr/rnnt/python/tools/export_parakeet_tdt_onnx.py
+++ b/asr/parakeet/model_export/export_parakeet.py
@@ -5,28 +5,19 @@
import argparse
import json
+import os
import shutil
import sys
from pathlib import Path
-from typing import Any
import torch
from torch import nn
-PYTHON_DIR = Path(__file__).resolve().parents[1]
-if str(PYTHON_DIR) not in sys.path:
- sys.path.insert(0, str(PYTHON_DIR))
+sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "common" / "model_export"))
+from nemo_export import configure_encoder_sdpa, export_nemo_preprocessor, load_nemo_model
-from rnnt.parakeet_tdt.nemo_backend import (
- DEFAULT_MODEL_ID,
- ONNX_DIR,
- apply_attention,
- configure_encoder_sdpa,
- load_nemo_model,
- write_tokenizer_artifacts,
-)
-
-from nemo_preprocessor_export import export_nemo_preprocessor
+DEFAULT_MODEL_ID = "nvidia/parakeet-tdt-0.6b-v3"
+ONNX_DIR = Path(os.environ.get("DIN_DEPLOY_ARTIFACTS_PATH", "artifacts")) / "parakeet" / "onnx"
ONNX_GRAPH_NAMES = ("preprocessor", "encoder", "predict_step")
@@ -89,12 +80,12 @@ def main() -> None:
model = load_nemo_model(args.model, device=device, dtype=dtype)
if args.local_window:
- apply_attention(model, "rel_pos_local_attn", [256, 256])
+ model.change_attention_model("rel_pos_local_attn", [256, 256])
+ model.to(device=device, dtype=dtype)
configure_encoder_sdpa(model.encoder)
extend_encoder_positional_encoding(model, args.encoder_profile_frames)
model.eval()
- model, quantization_info = restore_quantized_modelopt_checkpoint(model, args)
graphs = parse_onnx_graph_names(args.graphs)
if "preprocessor" in graphs:
export_preprocessor(model, args.output, args.opset, args.external_data, args.dynamo, args.optimize)
@@ -102,15 +93,7 @@ def main() -> None:
export_encoder(model, args.output, device, args.opset, args.external_data, args.dynamo, args.optimize)
if "predict_step" in graphs:
export_predict_step(model, args.output, device, args.opset, args.external_data, args.dynamo, args.optimize)
- if args.onnx_int4:
- quantization_info["onnx_int4"] = quantize_exported_onnx_graphs(
- args.output,
- args.onnx_int4_targets,
- args.external_data,
- args.onnx_int4_calibration_method,
- args.onnx_int4_calibration_eps,
- )
- write_metadata(model, args.output, args.dtype, args.encoder_profile_frames, quantization_info)
+ write_metadata(model, args.output, args.dtype, args.encoder_profile_frames)
write_tokenizer_artifacts(model, args.output)
@@ -194,69 +177,11 @@ def export_predict_step(
)
-def restore_quantized_modelopt_checkpoint(model, args: argparse.Namespace) -> tuple[object, dict[str, Any]]:
- if args.quantized_modelopt_path is None:
- return model, {}
- from rnnt.validation.quantization import restore_modelopt_quantized_model
-
- restored_model, info = restore_modelopt_quantized_model(
- model,
- args.quantized_modelopt_path,
- summary_path=args.output / "quant_summary.txt",
- )
- return restored_model, {"modelopt_checkpoint": info}
-
-
-def quantize_exported_onnx_graphs(
- output: Path,
- targets: str,
- external_data: bool,
- calibration_method: str,
- calibration_eps: str,
-) -> dict[str, Any]:
- try:
- import onnx
- from modelopt.onnx.quantization.int4 import quantize as quantize_int4 # type: ignore[import-not-found]
- except ImportError as error:
- raise RuntimeError(
- "ONNX INT4 export requires `pip install nvidia-modelopt` or `pip install -e .[quantization]`."
- ) from error
-
- graph_names = parse_onnx_graph_names(targets)
- providers = split_csv(calibration_eps)
- quantized_paths = []
- for graph_name in graph_names:
- onnx_path = output / f"{graph_name}.onnx"
- quantized_model = quantize_int4(
- str(onnx_path),
- calibration_method=calibration_method,
- calibration_data_reader=None,
- calibration_eps=providers,
- )
- onnx.save_model(
- quantized_model,
- onnx_path,
- save_as_external_data=external_data,
- location=onnx_path.name + "_data",
- size_threshold=0 if external_data else 1024,
- )
- quantized_paths.append(str(onnx_path))
- return {
- "mode": "modelopt-onnx-int4",
- "targets": graph_names,
- "calibration_method": calibration_method,
- "calibration_data": "modelopt-random",
- "calibration_eps": providers,
- "paths": quantized_paths,
- }
-
-
def write_metadata(
model,
output: Path,
dtype: str,
encoder_profile_frames: int,
- quantization_info: dict[str, Any] | None = None,
) -> None:
metadata = {
"blank_token_id": blank_token_id(model),
@@ -274,8 +199,6 @@ def write_metadata(
"onnx_dtype": dtype,
"boundary_dtype": "float32",
}
- if quantization_info:
- metadata["quantization"] = quantization_info
(output / "metadata.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8")
@@ -347,36 +270,24 @@ def extend_encoder_positional_encoding(model: nn.Module, encoder_profile_frames:
def validate_args(args: argparse.Namespace) -> None:
- parse_onnx_graph_names(args.onnx_int4_targets)
- parse_onnx_graph_names(args.graphs)
- if args.encoder_profile_frames < 16:
- raise ValueError("--encoder-profile-frames must be at least 16.")
- if args.onnx_int4 and args.quantized_modelopt_path is not None:
- raise ValueError("Choose either --onnx-int4 or --quantized-modelopt-path for INT4 export, not both.")
- if (args.onnx_int4 or args.quantized_modelopt_path is not None) and args.opset < 21:
- raise ValueError("INT4 ONNX export requires --opset 21 or newer.")
+ if args.encoder_profile_frames <= 0:
+ raise ValueError("--encoder-profile-frames must be positive.")
def parse_onnx_graph_names(value: str) -> list[str]:
- names = split_csv(value)
+ names = [name.strip() for name in value.split(",") if name.strip()]
if not names:
- raise ValueError("At least one ONNX INT4 target is required.")
+ raise ValueError("At least one ONNX graph is required.")
if "all" in names:
if len(names) > 1:
- raise ValueError("--onnx-int4-targets cannot combine 'all' with explicit graph names.")
+ raise ValueError("--graphs cannot combine 'all' with explicit graph names.")
return list(ONNX_GRAPH_NAMES)
unknown = sorted(set(names) - set(ONNX_GRAPH_NAMES))
if unknown:
- raise ValueError(f"Unsupported ONNX INT4 target(s): {', '.join(unknown)}.")
+ raise ValueError(f"Unsupported ONNX graph(s): {', '.join(unknown)}.")
return names
-def split_csv(value: str | None) -> list[str]:
- if not value:
- return []
- return [part.strip() for part in value.split(",") if part.strip()]
-
-
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Export Parakeet TDT ONNX artifacts for the C++ baseline.")
parser.add_argument("--model", default=DEFAULT_MODEL_ID)
@@ -396,29 +307,6 @@ def parse_args() -> argparse.Namespace:
default=65536,
help="Maximum mel frames used by the C++ TensorRT encoder profile.",
)
- parser.add_argument(
- "--quantized-modelopt-path",
- type=Path,
- default=None,
- help="Restore a saved ModelOpt checkpoint before exporting with torch.onnx.export.",
- )
- parser.add_argument(
- "--onnx-int4",
- action=argparse.BooleanOptionalAction,
- default=False,
- help="Post-process exported ONNX graphs with ModelOpt ONNX INT4 PTQ. Requires --opset 21 or newer.",
- )
- parser.add_argument(
- "--onnx-int4-targets",
- default="encoder",
- help="Comma-separated ONNX graphs to quantize: preprocessor, encoder, predict_step, or all.",
- )
- parser.add_argument("--onnx-int4-calibration-method", default="awq_lite")
- parser.add_argument(
- "--onnx-int4-calibration-eps",
- default="dml,cpu",
- help="Comma-separated calibration execution providers passed to ModelOpt ONNX INT4 PTQ.",
- )
return parser.parse_args()
@@ -447,6 +335,30 @@ def joint_from_projected_encoder(
return logits
+def write_tokenizer_artifacts(model: nn.Module, output: Path) -> None:
+ tokenizer = getattr(model, "tokenizer", None)
+ if tokenizer is None:
+ raise RuntimeError("NeMo Parakeet model does not expose a tokenizer.")
+
+ sentencepiece = getattr(tokenizer, "tokenizer", tokenizer)
+ if hasattr(sentencepiece, "get_piece_size") and hasattr(sentencepiece, "id_to_piece"):
+ vocab = {sentencepiece.id_to_piece(index): index for index in range(sentencepiece.get_piece_size())}
+ (output / "tokenizer.json").write_text(json.dumps({"model": {"vocab": vocab}}, indent=2), encoding="utf-8")
+ return
+ if hasattr(tokenizer, "tokenizer") and hasattr(tokenizer.tokenizer, "save"):
+ tokenizer.tokenizer.save(str(output / "tokenizer.json"))
+ return
+ if hasattr(tokenizer, "save_to"):
+ tokenizer.save_to(str(output / "tokenizer.json"))
+ return
+ if hasattr(tokenizer, "vocab_path"):
+ vocab_path = Path(tokenizer.vocab_path)
+ if vocab_path.is_file():
+ (output / "tokenizer.json").write_bytes(vocab_path.read_bytes())
+ return
+ raise RuntimeError("Unable to write tokenizer.json from NeMo Parakeet tokenizer.")
+
+
if __name__ == "__main__":
with torch.inference_mode():
main()
diff --git a/asr/parakeet/model_export/validate_parakeet.py b/asr/parakeet/model_export/validate_parakeet.py
new file mode 100644
index 0000000..edf5975
--- /dev/null
+++ b/asr/parakeet/model_export/validate_parakeet.py
@@ -0,0 +1,24 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "common" / "model_export"))
+from export_parakeet import DEFAULT_MODEL_ID
+from nemo_export import load_nemo_model
+from validate_asr import validate, validation_parser
+
+
+def reference(args):
+ model = load_nemo_model(args.model, device=args.device)
+ if args.local_window:
+ model.change_attention_model("rel_pos_local_attn", [256, 256])
+ result = model.transcribe([str(args.audio.resolve())], batch_size=1)[0]
+ return str(getattr(result, "text", result)).strip()
+
+
+if __name__ == "__main__":
+ parser = validation_parser(DEFAULT_MODEL_ID)
+ parser.add_argument("--local-window", action="store_true", help="Match an export made with --local-window.")
+ validate(parser.parse_args(), reference)
diff --git a/asr/rnnt/cpp/parakeet_tdt.cpp b/asr/parakeet/parakeet_tdt.cpp
similarity index 100%
rename from asr/rnnt/cpp/parakeet_tdt.cpp
rename to asr/parakeet/parakeet_tdt.cpp
diff --git a/asr/rnnt/cpp/parakeet_tdt.h b/asr/parakeet/parakeet_tdt.h
similarity index 100%
rename from asr/rnnt/cpp/parakeet_tdt.h
rename to asr/parakeet/parakeet_tdt.h
diff --git a/asr/parakeet/requirements.txt b/asr/parakeet/requirements.txt
new file mode 100644
index 0000000..eb67fa9
--- /dev/null
+++ b/asr/parakeet/requirements.txt
@@ -0,0 +1 @@
+-r ../../common/model_export/requirements-nemo.txt
diff --git a/asr/rnnt/CMakeLists.txt b/asr/rnnt/CMakeLists.txt
deleted file mode 100644
index 56a2920..0000000
--- a/asr/rnnt/CMakeLists.txt
+++ /dev/null
@@ -1,134 +0,0 @@
-option(DIN_ENABLE_PARAKEET_PYTHON_BINDINGS "Build Parakeet TDT nanobind Python bindings" OFF)
-option(DIN_ENABLE_NEMOTRON_PYTHON_BINDINGS "Build Nemotron ASR nanobind Python bindings" OFF)
-
-add_din_shared_library(din_asr_parakeet_tdt SHARED
- cpp/parakeet_tdt.cpp
-)
-
-target_include_directories(din_asr_parakeet_tdt
- PUBLIC
- "${CMAKE_CURRENT_SOURCE_DIR}/cpp"
-)
-target_link_libraries(din_asr_parakeet_tdt
- PUBLIC
- din_common_ort
- din_common_io
- nlohmann_json::nlohmann_json
- argparse
-)
-
-set_target_properties(din_asr_parakeet_tdt PROPERTIES
- BUILD_WITH_INSTALL_RPATH ON
- INSTALL_RPATH "$ORIGIN"
- WINDOWS_EXPORT_ALL_SYMBOLS ON
-)
-
-add_din_executable(din_asr_parakeet_tdt_cli
- cpp/parakeet_tdt_main.cpp
-)
-target_link_libraries(din_asr_parakeet_tdt_cli PRIVATE din_asr_parakeet_tdt)
-
-add_din_shared_library(din_asr_nemotron SHARED
- cpp/nemotron.cpp
-)
-
-target_include_directories(din_asr_nemotron
- PUBLIC
- "${CMAKE_CURRENT_SOURCE_DIR}/cpp"
-)
-target_link_libraries(din_asr_nemotron
- PUBLIC
- din_common_ort
- din_common_io
- nlohmann_json::nlohmann_json
- argparse
-)
-
-set_target_properties(din_asr_nemotron PROPERTIES
- BUILD_WITH_INSTALL_RPATH ON
- INSTALL_RPATH "$ORIGIN"
- WINDOWS_EXPORT_ALL_SYMBOLS ON
-)
-
-add_din_executable(din_asr_nemotron_cli
- cpp/nemotron_main.cpp
-)
-target_link_libraries(din_asr_nemotron_cli PRIVATE din_asr_nemotron)
-
-if(DIN_ENABLE_PARAKEET_PYTHON_BINDINGS OR DIN_ENABLE_NEMOTRON_PYTHON_BINDINGS)
- find_package(Python REQUIRED COMPONENTS Interpreter Development.Module OPTIONAL_COMPONENTS Development.SABIModule)
- include(FetchContent)
- FetchContent_Declare(
- nanobind
- GIT_REPOSITORY https://github.com/wjakob/nanobind.git
- GIT_TAG v2.9.2
- GIT_SHALLOW TRUE
- )
- FetchContent_MakeAvailable(nanobind)
-endif()
-
-function(copy_rnnt_python_runtime target package_name)
- set(_rnnt_python_runtime_files)
- if(WIN32)
- if(ONNXRUNTIME_DLL)
- list(APPEND _rnnt_python_runtime_files "${ONNXRUNTIME_DLL}")
- endif()
- if(ONNXRUNTIME_PROVIDERS_SHARED_DLL)
- list(APPEND _rnnt_python_runtime_files "${ONNXRUNTIME_PROVIDERS_SHARED_DLL}")
- endif()
- else()
- if(ONNXRUNTIME_LIB)
- list(APPEND _rnnt_python_runtime_files "${ONNXRUNTIME_LIB}")
- get_filename_component(_onnxruntime_lib_dir "${ONNXRUNTIME_LIB}" DIRECTORY)
- file(GLOB _onnxruntime_soname_files "${_onnxruntime_lib_dir}/libonnxruntime.so.*")
- list(APPEND _rnnt_python_runtime_files ${_onnxruntime_soname_files})
- endif()
- if(ONNXRUNTIME_PROVIDERS_SHARED_LIB)
- list(APPEND _rnnt_python_runtime_files "${ONNXRUNTIME_PROVIDERS_SHARED_LIB}")
- endif()
- endif()
-
- if(_rnnt_python_runtime_files)
- add_custom_command(TARGET ${target} POST_BUILD
- COMMAND "${CMAKE_COMMAND}" -E copy_if_different
- ${_rnnt_python_runtime_files}
- "$"
- COMMAND_EXPAND_LISTS
- )
- install(FILES ${_rnnt_python_runtime_files} DESTINATION ${package_name})
- endif()
-endfunction()
-
-if(DIN_ENABLE_PARAKEET_PYTHON_BINDINGS)
- nanobind_add_module(_parakeet_tdt_cpp
- STABLE_ABI
- NB_STATIC
- cpp/parakeet_tdt_python_bindings.cpp
- )
- target_link_libraries(_parakeet_tdt_cpp PRIVATE din_asr_parakeet_tdt)
- set_target_properties(_parakeet_tdt_cpp PROPERTIES
- BUILD_WITH_INSTALL_RPATH ON
- BUILD_RPATH "$ORIGIN"
- INSTALL_RPATH "$ORIGIN"
- )
- install(TARGETS _parakeet_tdt_cpp LIBRARY DESTINATION rnnt RUNTIME DESTINATION rnnt)
- install(TARGETS din_asr_parakeet_tdt LIBRARY DESTINATION rnnt RUNTIME DESTINATION rnnt)
- copy_rnnt_python_runtime(_parakeet_tdt_cpp rnnt)
-endif()
-
-if(DIN_ENABLE_NEMOTRON_PYTHON_BINDINGS)
- nanobind_add_module(_nemotron_asr_cpp
- STABLE_ABI
- NB_STATIC
- cpp/nemotron_python_bindings.cpp
- )
- target_link_libraries(_nemotron_asr_cpp PRIVATE din_asr_nemotron)
- set_target_properties(_nemotron_asr_cpp PROPERTIES
- BUILD_WITH_INSTALL_RPATH ON
- BUILD_RPATH "$ORIGIN"
- INSTALL_RPATH "$ORIGIN"
- )
- install(TARGETS _nemotron_asr_cpp LIBRARY DESTINATION rnnt RUNTIME DESTINATION rnnt)
- install(TARGETS din_asr_nemotron LIBRARY DESTINATION rnnt RUNTIME DESTINATION rnnt)
- copy_rnnt_python_runtime(_nemotron_asr_cpp rnnt)
-endif()
diff --git a/asr/rnnt/README.md b/asr/rnnt/README.md
deleted file mode 100644
index fe0ab64..0000000
--- a/asr/rnnt/README.md
+++ /dev/null
@@ -1,32 +0,0 @@
-# RNNT ASR
-
-ONNX transcription pipelines for Parakeet TDT and Nemotron ASR.
-
-## Supported models
-
-| Model | Hugging Face ID | Checkpoint | ONNX export |
-| --- | --- | --- | --- |
-| Parakeet TDT 0.6B v3 | `nvidia/parakeet-tdt-0.6b-v3` | FP32 | FP16 (default) |
-| Nemotron 3.5 ASR Streaming 0.6B | `nvidia/nemotron-3.5-asr-streaming-0.6b` | FP32 | FP16 (default) |
-
-## Export
-
-```bash
-PYTHONPATH=asr/rnnt/python python asr/rnnt/python/tools/export_parakeet_tdt_onnx.py --help
-PYTHONPATH=asr/rnnt/python python asr/rnnt/python/tools/export_nemotron_onnx.py --help
-```
-
-## Build
-
-```bash
-cmake --preset linux-x64 -DDIN_ENABLE_PARAKEET_PYTHON_BINDINGS=ON -DDIN_ENABLE_NEMOTRON_PYTHON_BINDINGS=ON
-cmake --build out/build/linux-x64 --target din_asr_parakeet_tdt_cli din_asr_nemotron_cli _parakeet_tdt_cpp _nemotron_asr_cpp
-```
-
-## Run
-
-```bash
-out/build/linux-x64/bin/din_asr_parakeet_tdt_cli audio.wav --model-dir artifacts/parakeet/onnx --provider trt-rtx
-out/build/linux-x64/bin/din_asr_nemotron_cli audio.wav --model-dir artifacts/nemotron/onnx --provider trt-rtx --lang-id en-US
-```
-
diff --git a/asr/rnnt/cpp/asr_bindings.h b/asr/rnnt/cpp/asr_bindings.h
deleted file mode 100644
index 8f4aa44..0000000
--- a/asr/rnnt/cpp/asr_bindings.h
+++ /dev/null
@@ -1,92 +0,0 @@
-// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-// SPDX-License-Identifier: Apache-2.0
-
-#pragma once
-
-#include
-#include
-#include
-
-#include "asr_common.h"
-#include "audio.h"
-#include "ort_helper.h"
-#include
-#include
-#include
-#include
-
-namespace din::asr::common
-{
-
-namespace nb = nanobind;
-
-struct BindingDefaults
-{
- const char* model_dir;
- const char* provider;
- const char* ep_cache;
- const char* ep_context;
-};
-
-// Generic nanobind transcriber. The per-model result is converted to a dict
-// through an unqualified ToDict() resolved by argument-dependent lookup in the
-// model's namespace, which is the only model-specific piece of the binding.
-template
-class Transcriber
-{
-public:
- Transcriber(const std::string& model_dir, const std::string& provider, const std::string& ep_cache,
- const std::string& ep_context_dir, const std::string& lang_id)
- {
- Config config;
- config.model_dir = model_dir;
- config.provider = provider;
- config.ep_cache_dir = ep_cache;
- config.ep_context_dir = ep_context_dir;
- if constexpr (requires(Config value) { value.lang_id = std::string{}; })
- {
- config.lang_id = lang_id;
- }
- pipeline_ = std::make_unique(std::move(config));
- }
-
- nb::dict Transcribe(nb::ndarray, nb::c_contig> samples, int sample_rate)
- {
- din::io::Audio audio;
- audio.sample_rate = sample_rate;
- audio.samples.assign(samples.data(), samples.data() + samples.shape(0));
-
- nb::gil_scoped_release release;
- const auto result = pipeline_->Transcribe(audio);
- nb::gil_scoped_acquire acquire;
- return ToDict(result);
- }
-
- nb::dict TranscribeFile(const std::string& path)
- {
- nb::gil_scoped_release release;
- const auto result = pipeline_->TranscribeFile(path);
- nb::gil_scoped_acquire acquire;
- return ToDict(result);
- }
-
-private:
- std::unique_ptr pipeline_;
-};
-
-template
-void BindTranscriber(nb::module_& module, const char* doc, const BindingDefaults& defaults)
-{
- using Bound = Transcriber;
- module.doc() = doc;
- nb::class_(module, "Transcriber")
- .def(nb::init(),
- nb::arg("model_dir") = defaults.model_dir, nb::arg("provider") = defaults.provider,
- nb::arg("ep_cache") = defaults.ep_cache, nb::arg("ep_context_dir") = defaults.ep_context,
- nb::arg("lang_id") = "auto")
- .def("transcribe", &Bound::Transcribe, nb::arg("samples"), nb::arg("sample_rate"))
- .def("transcribe_file", &Bound::TranscribeFile, nb::arg("path"));
-}
-
-} // namespace din::asr::common
diff --git a/asr/rnnt/cpp/nemotron_python_bindings.cpp b/asr/rnnt/cpp/nemotron_python_bindings.cpp
deleted file mode 100644
index a2ea00f..0000000
--- a/asr/rnnt/cpp/nemotron_python_bindings.cpp
+++ /dev/null
@@ -1,31 +0,0 @@
-// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-// SPDX-License-Identifier: Apache-2.0
-
-#include "asr_bindings.h"
-#include "nemotron.h"
-
-namespace din::asr::nemotron
-{
-
-namespace nb = nanobind;
-
-nb::dict ToDict(const TranscriptionResult& result)
-{
- nb::dict output;
- output["text"] = result.text;
- output["tokens"] = result.generated.tokens;
- output["starts"] = result.generated.starts;
- return output;
-}
-
-} // namespace din::asr::nemotron
-
-NB_MODULE(_nemotron_asr_cpp, module)
-{
- din::asr::common::BindTranscriber(
- module, "nanobind wrapper for the DIN Deploy Nemotron C++ ONNX pipeline",
- {.model_dir = "artifacts/nemotron/onnx",
- .provider = "cpu",
- .ep_cache = "artifacts/nemotron/trt_rtx_cache",
- .ep_context = "artifacts/nemotron/ep_context"});
-}
diff --git a/asr/rnnt/cpp/parakeet_tdt_python_bindings.cpp b/asr/rnnt/cpp/parakeet_tdt_python_bindings.cpp
deleted file mode 100644
index 90f5a5a..0000000
--- a/asr/rnnt/cpp/parakeet_tdt_python_bindings.cpp
+++ /dev/null
@@ -1,32 +0,0 @@
-// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-// SPDX-License-Identifier: Apache-2.0
-
-#include "asr_bindings.h"
-#include "parakeet_tdt.h"
-
-namespace din::asr::parakeet
-{
-
-namespace nb = nanobind;
-
-nb::dict ToDict(const TranscriptionResult& result)
-{
- nb::dict output;
- output["text"] = result.text;
- output["tokens"] = result.generated.tokens;
- output["durations"] = result.generated.durations;
- output["starts"] = result.generated.starts;
- return output;
-}
-
-} // namespace din::asr::parakeet
-
-NB_MODULE(_parakeet_tdt_cpp, module)
-{
- din::asr::common::BindTranscriber(
- module, "nanobind wrapper for the DIN Deploy Parakeet TDT C++ ONNX pipeline",
- {.model_dir = "artifacts/parakeet/onnx",
- .provider = "trt-rtx",
- .ep_cache = "artifacts/parakeet/trt_rtx_cache",
- .ep_context = "artifacts/parakeet/ep_context"});
-}
diff --git a/asr/rnnt/python/rnnt/__init__.py b/asr/rnnt/python/rnnt/__init__.py
deleted file mode 100644
index 52a7a9d..0000000
--- a/asr/rnnt/python/rnnt/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
diff --git a/asr/rnnt/python/rnnt/audio.py b/asr/rnnt/python/rnnt/audio.py
deleted file mode 100644
index 0949377..0000000
--- a/asr/rnnt/python/rnnt/audio.py
+++ /dev/null
@@ -1,72 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-from pathlib import Path
-
-import numpy as np
-import soundfile as sf
-
-
-@dataclass(frozen=True)
-class Audio:
- samples: np.ndarray
- sampling_rate: int
-
- @property
- def duration(self) -> float:
- return float(len(self.samples) / self.sampling_rate)
-
-
-def audio_duration(path: str | Path) -> float:
- info = sf.info(str(path))
- return float(info.frames / info.samplerate)
-
-
-def load_audio(path: str | Path, target_sampling_rate: int) -> Audio:
- samples, sampling_rate = sf.read(str(path), dtype="float32", always_2d=False)
- if samples.ndim == 2:
- samples = samples.mean(axis=1)
- if sampling_rate != target_sampling_rate:
- samples = resample(samples, sampling_rate, target_sampling_rate)
- sampling_rate = target_sampling_rate
- return Audio(samples=np.asarray(samples, dtype=np.float32), sampling_rate=sampling_rate)
-
-
-def resample(samples: np.ndarray, source_rate: int, target_rate: int) -> np.ndarray:
- if source_rate <= 0 or target_rate <= 0:
- raise ValueError("Sampling rates must be positive.")
- if len(samples) == 0 or source_rate == target_rate:
- return samples.astype(np.float32, copy=False)
- duration = len(samples) / source_rate
- source_times = np.linspace(0.0, duration, num=len(samples), endpoint=False)
- target_len = max(1, round(duration * target_rate))
- target_times = np.linspace(0.0, duration, num=target_len, endpoint=False)
- return np.interp(target_times, source_times, samples).astype(np.float32)
-
-
-def iter_audio_windows(
- audio: Audio,
- window_seconds: float,
- stride_seconds: float | None = None,
-):
- if window_seconds <= 0:
- raise ValueError("window_seconds must be positive.")
-
- stride_seconds = stride_seconds or window_seconds
- if stride_seconds <= 0:
- raise ValueError("stride_seconds must be positive.")
-
- window = max(1, round(window_seconds * audio.sampling_rate))
- stride = max(1, round(stride_seconds * audio.sampling_rate))
-
- start = 0
- while start < len(audio.samples):
- end = min(start + window, len(audio.samples))
- offset = start / audio.sampling_rate
- yield offset, Audio(audio.samples[start:end], audio.sampling_rate)
- if end == len(audio.samples):
- break
- start += stride
diff --git a/asr/rnnt/python/rnnt/nemotron_asr/__init__.py b/asr/rnnt/python/rnnt/nemotron_asr/__init__.py
deleted file mode 100644
index db7ee47..0000000
--- a/asr/rnnt/python/rnnt/nemotron_asr/__init__.py
+++ /dev/null
@@ -1,16 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-__all__ = [
- "DEFAULT_MODEL_ID",
- "load_nemo_model",
- "transcribe_paths",
-]
-
-
-def __getattr__(name: str):
- if name in {"DEFAULT_MODEL_ID", "load_nemo_model", "transcribe_paths"}:
- from . import nemo_backend
-
- return getattr(nemo_backend, name)
- raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
diff --git a/asr/rnnt/python/rnnt/nemotron_asr/cli.py b/asr/rnnt/python/rnnt/nemotron_asr/cli.py
deleted file mode 100644
index 49d68bb..0000000
--- a/asr/rnnt/python/rnnt/nemotron_asr/cli.py
+++ /dev/null
@@ -1,77 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-import argparse
-import time
-from pathlib import Path
-
-
-def main() -> None:
- from rnnt.audio import audio_duration
-
- from .nemo_backend import DEFAULT_ATT_CONTEXT_SIZE, load_nemo_model, transcribe_paths
-
- args = parse_args()
- started = time.perf_counter()
- duration_seconds = sum(audio_duration(path) for path in args.audiofile)
-
- load_started = time.perf_counter()
- model = load_nemo_model(
- args.model,
- device=args.device,
- target_lang=args.target_lang,
- att_context_size=DEFAULT_ATT_CONTEXT_SIZE,
- )
- load_seconds = time.perf_counter() - load_started
-
- transcribe_started = time.perf_counter()
- texts = transcribe_paths(model, args.audiofile, args.target_lang)
- transcribe_seconds = time.perf_counter() - transcribe_started
-
- for path, text in zip(args.audiofile, texts, strict=True):
- if args.strip_lang_tags:
- from .nemo_backend import strip_lang_tags
-
- text = strip_lang_tags(text)
- if len(args.audiofile) > 1:
- print(f"{path}: {text}")
- else:
- print(text)
-
- total_seconds = time.perf_counter() - started
- print()
- print(f"device: {model.device}")
- print(f"audio: {duration_seconds:.2f}s")
- print(f"load: {load_seconds:.2f}s")
- print(f"transcribe: {transcribe_seconds:.2f}s")
- print(f"rtf: {transcribe_seconds / duration_seconds:.3f}")
- print(f"total: {total_seconds:.2f}s")
-
-
-def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
- parser = argparse.ArgumentParser(description="Transcribe audio with Nemotron 3.5 ASR.")
- parser.add_argument("audiofile", nargs="+", type=Path)
- parser.add_argument("--model", default="nvidia/nemotron-3.5-asr-streaming-0.6b")
- parser.add_argument("--device", default="cuda")
- parser.add_argument("--target-lang", "--target_lang", default="auto")
- parser.add_argument("--strip-lang-tags", "--strip_lang_tags", type=parse_bool, nargs="?", const=True, default=True)
- parser.add_argument("--no-strip-lang-tags", "--no_strip_lang_tags", dest="strip_lang_tags", action="store_false")
- parser.add_argument("--keep-lang-tags", dest="strip_lang_tags", action="store_false")
- return parser.parse_args(argv)
-
-
-def parse_bool(value: str | bool) -> bool:
- if isinstance(value, bool):
- return value
- normalized = value.strip().lower()
- if normalized in {"1", "true", "t", "yes", "y", "on"}:
- return True
- if normalized in {"0", "false", "f", "no", "n", "off"}:
- return False
- raise argparse.ArgumentTypeError("boolean value must be true or false.")
-
-
-if __name__ == "__main__":
- main()
diff --git a/asr/rnnt/python/rnnt/nemotron_asr/cpp_backend.py b/asr/rnnt/python/rnnt/nemotron_asr/cpp_backend.py
deleted file mode 100644
index db1851b..0000000
--- a/asr/rnnt/python/rnnt/nemotron_asr/cpp_backend.py
+++ /dev/null
@@ -1,48 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-from pathlib import Path
-from typing import Any
-
-import numpy as np
-
-from rnnt.audio import Audio
-
-
-class CppBackend:
- def __init__(
- self,
- *,
- model_dir: str | Path = Path("artifacts/nemotron/onnx"),
- provider: str = "cpu",
- ep_cache: str | Path = Path("artifacts/nemotron/trt_rtx_cache"),
- ep_context_dir: str | Path = Path("artifacts/nemotron/ep_context"),
- lang_id: str = "auto",
- ) -> None:
- try:
- from rnnt._nemotron_asr_cpp import Transcriber
- except ImportError as error:
- try:
- from _nemotron_asr_cpp import Transcriber
- except ImportError:
- raise RuntimeError(
- "Nemotron C++ bindings are not importable. Build `_nemotron_asr_cpp` "
- "with CMake and add the build output directory to PYTHONPATH."
- ) from error
-
- self.transcriber = Transcriber(
- str(model_dir),
- provider,
- str(ep_cache),
- str(ep_context_dir),
- lang_id,
- )
-
- def transcribe_audio(self, audio: Audio) -> dict[str, Any]:
- samples = np.ascontiguousarray(audio.samples, dtype=np.float32)
- return dict(self.transcriber.transcribe(samples, int(audio.sampling_rate)))
-
- def transcribe_file(self, path: str | Path) -> dict[str, Any]:
- return dict(self.transcriber.transcribe_file(str(path)))
diff --git a/asr/rnnt/python/rnnt/nemotron_asr/nemo_backend.py b/asr/rnnt/python/rnnt/nemotron_asr/nemo_backend.py
deleted file mode 100644
index 4e3faff..0000000
--- a/asr/rnnt/python/rnnt/nemotron_asr/nemo_backend.py
+++ /dev/null
@@ -1,154 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-import re
-import sys
-import tempfile
-from pathlib import Path
-from typing import Any
-
-import soundfile as sf
-import torch
-from torch import nn
-
-from rnnt.audio import Audio
-
-DEFAULT_MODEL_ID = "nvidia/nemotron-3.5-asr-streaming-0.6b"
-DEFAULT_ATT_CONTEXT_SIZE = (56, 3)
-LANG_TAG_RE = re.compile(r"<[^>\s]+>")
-
-
-def strip_lang_tags(text: str) -> str:
- return LANG_TAG_RE.sub("", text).strip()
-
-
-def load_nemo_model(
- model_id: str = DEFAULT_MODEL_ID,
- *,
- device: str | torch.device | None = None,
- target_lang: str = "auto",
- att_context_size: tuple[int, int] = DEFAULT_ATT_CONTEXT_SIZE,
- dtype: torch.dtype = torch.float32,
- configure_for_export: bool = False,
-) -> nn.Module:
- target_device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
- nemo_model = _load_nemo_asr_model(model_id, target_device)
- nemo_model = nemo_model.to(device=target_device, dtype=dtype)
- nemo_model.eval()
- _configure_nemo_model(
- nemo_model,
- target_lang=target_lang,
- att_context_size=att_context_size,
- configure_for_export=configure_for_export,
- )
- return nemo_model
-
-
-def prompt_dictionary(nemo_model: nn.Module) -> dict[str, int]:
- for key in ("train_ds", "validation_ds", "test_ds"):
- data = nemo_model.cfg.get(key)
- if data is not None and "prompt_dictionary" in data:
- return {str(lang): int(prompt_id) for lang, prompt_id in data.prompt_dictionary.items()}
- raise KeyError("prompt_dictionary")
-
-
-def prompt_id(nemo_model: nn.Module, target_lang: str) -> int:
- prompts = prompt_dictionary(nemo_model)
- if target_lang not in prompts:
- raise ValueError(f"Unsupported target language: {target_lang}")
- return prompts[target_lang]
-
-
-def transcribe_paths(nemo_model: nn.Module, paths: list[str | Path], target_lang: str) -> list[str]:
- override_config = _prompt_transcribe_config(target_lang)
- results = nemo_model.transcribe([str(path) for path in paths], override_config=override_config)
- return [getattr(result, "text", str(result)).strip() for result in results]
-
-
-def transcribe_audio(nemo_model: nn.Module, audio: Audio, target_lang: str) -> str:
- with tempfile.NamedTemporaryFile(suffix=".wav") as handle:
- sf.write(handle.name, audio.samples, audio.sampling_rate)
- return transcribe_paths(nemo_model, [handle.name], target_lang)[0]
-
-
-def write_tokenizer_artifacts(nemo_model: nn.Module, output: Path) -> None:
- tokenizer = nemo_model.tokenizer.tokenizer
- (output / "tokenizer.model").write_bytes(tokenizer.serialized_model_proto())
- pieces = [tokenizer.id_to_piece(index) for index in range(tokenizer.get_piece_size())]
- (output / "vocab.txt").write_text("\n".join(pieces) + "\n", encoding="utf-8")
-
-
-def configure_nemo_encoder_sdpa(encoder: nn.Module) -> None:
- encoder_flag_updated = False
- if hasattr(encoder, "use_pytorch_sdpa"):
- encoder.use_pytorch_sdpa = True
- encoder_flag_updated = True
-
- layer_flags_updated = 0
- attention_flags_updated = 0
- attention_classes: dict[str, int] = {}
- for layer in getattr(encoder, "layers", []):
- if hasattr(layer, "use_pytorch_sdpa"):
- layer.use_pytorch_sdpa = True
- layer_flags_updated += 1
- self_attn = getattr(layer, "self_attn", None)
- if self_attn is None:
- continue
- attention_classes[type(self_attn).__name__] = attention_classes.get(type(self_attn).__name__, 0) + 1
- if hasattr(self_attn, "use_pytorch_sdpa"):
- self_attn.use_pytorch_sdpa = True
- attention_flags_updated += 1
-
- print(
- "NeMo encoder SDPA export config: "
- f"encoder_flag={encoder_flag_updated}, "
- f"layer_flags={layer_flags_updated}, "
- f"attention_flags={attention_flags_updated}, "
- f"attention_classes={attention_classes}",
- file=sys.stderr,
- )
-
-
-def _load_nemo_asr_model(model_id: str, device: torch.device) -> nn.Module:
- try:
- from nemo.collections import asr as nemo_asr
- except ImportError as exc:
- raise RuntimeError("Nemotron requires the optional `nemo_toolkit[asr]` dependency.") from exc
-
- model_path = Path(model_id)
- if model_path.is_file() and model_path.suffix == ".nemo":
- return nemo_asr.models.ASRModel.restore_from(str(model_path), map_location=device)
- return nemo_asr.models.ASRModel.from_pretrained(model_name=model_id, map_location=device)
-
-
-def _configure_nemo_model(
- nemo_model: nn.Module,
- *,
- target_lang: str,
- att_context_size: tuple[int, int],
- configure_for_export: bool,
-) -> None:
- if hasattr(nemo_model, "set_inference_prompt"):
- nemo_model.set_inference_prompt(target_lang)
- if configure_for_export and hasattr(nemo_model, "set_export_config"):
- nemo_model.set_export_config({"cache_support": True})
- encoder = nemo_model.encoder
- if hasattr(encoder, "set_default_att_context_size"):
- encoder.set_default_att_context_size(list(att_context_size))
- if configure_for_export:
- configure_nemo_encoder_sdpa(encoder)
- if hasattr(encoder, "export_cache_support"):
- encoder.export_cache_support = True
- encoder.eval()
-
-
-def _prompt_transcribe_config(target_lang: str) -> Any | None:
- try:
- from nemo.collections.asr.models.rnnt_bpe_models_prompt import RNNTPromptTranscribeConfig
- except Exception:
- return None
- override_config = RNNTPromptTranscribeConfig(batch_size=1, verbose=False, target_lang=target_lang)
- override_config.use_lhotse = False
- return override_config
diff --git a/asr/rnnt/python/rnnt/parakeet_tdt/__init__.py b/asr/rnnt/python/rnnt/parakeet_tdt/__init__.py
deleted file mode 100644
index 1cfab50..0000000
--- a/asr/rnnt/python/rnnt/parakeet_tdt/__init__.py
+++ /dev/null
@@ -1,12 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-__all__ = ["load_nemo_model", "transcribe_audio", "transcribe_paths"]
-
-
-def __getattr__(name: str):
- if name in __all__:
- from . import nemo_backend
-
- return getattr(nemo_backend, name)
- raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
diff --git a/asr/rnnt/python/rnnt/parakeet_tdt/cli.py b/asr/rnnt/python/rnnt/parakeet_tdt/cli.py
deleted file mode 100644
index 12612dd..0000000
--- a/asr/rnnt/python/rnnt/parakeet_tdt/cli.py
+++ /dev/null
@@ -1,63 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-import argparse
-import time
-from pathlib import Path
-
-import soundfile as sf
-import torch
-
-from .nemo_backend import apply_attention, load_nemo_model, transcribe_paths
-
-
-def main() -> None:
- args = parse_args()
- started = time.perf_counter()
- duration_seconds = audio_duration(args.audiofile)
-
- load_started = time.perf_counter()
- model = load_nemo_model(args.model, device=args.device)
- load_seconds = time.perf_counter() - load_started
-
- if args.local_window:
- apply_attention(model, "rel_pos_local_attn", [256, 256])
-
- transcribe_started = time.perf_counter()
- if args.timestamps != "none":
- raise ValueError("Parakeet NeMo CLI currently uses NeMo text transcription; token timestamps are not exposed.")
- text = transcribe_paths(model, [args.audiofile])[0]
- transcribe_seconds = time.perf_counter() - transcribe_started
-
- print(text)
-
- total_seconds = time.perf_counter() - started
- print()
- print(f"device: {args.device}")
- print(f"audio: {duration_seconds:.2f}s")
- print(f"load: {load_seconds:.2f}s")
- print(f"transcribe: {transcribe_seconds:.2f}s")
- print(f"rtf: {transcribe_seconds / duration_seconds:.3f}")
- print(f"total: {total_seconds:.2f}s")
-
-
-def parse_args() -> argparse.Namespace:
- parser = argparse.ArgumentParser(description="Transcribe audio with NeMo Parakeet TDT.")
- parser.add_argument("audiofile", type=Path)
- parser.add_argument("--model", default="nvidia/parakeet-tdt-0.6b-v3")
- parser.add_argument("--device", default="cuda")
- parser.add_argument("--timestamps", "--timesteps", choices=("none", "segment", "token"), default="none")
- parser.add_argument("--local-window", action=argparse.BooleanOptionalAction, default=False)
- return parser.parse_args()
-
-
-def audio_duration(path: Path) -> float:
- info = sf.info(str(path))
- return float(info.frames / info.samplerate)
-
-
-if __name__ == "__main__":
- with torch.inference_mode():
- main()
diff --git a/asr/rnnt/python/rnnt/parakeet_tdt/cpp_backend.py b/asr/rnnt/python/rnnt/parakeet_tdt/cpp_backend.py
deleted file mode 100644
index 786598c..0000000
--- a/asr/rnnt/python/rnnt/parakeet_tdt/cpp_backend.py
+++ /dev/null
@@ -1,47 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-from pathlib import Path
-from typing import Any
-
-import numpy as np
-
-from rnnt.audio import Audio
-
-
-class CppBackend:
- def __init__(
- self,
- *,
- model_dir: str | Path = Path("artifacts/parakeet/onnx"),
- provider: str = "cpu",
- ep_cache: str | Path = Path("artifacts/parakeet/trt_rtx_cache"),
- ep_context_dir: str | Path = Path("artifacts/parakeet/ep_context"),
- ) -> None:
- try:
- from rnnt._parakeet_tdt_cpp import Transcriber
- except ImportError as error:
- try:
- from _parakeet_tdt_cpp import Transcriber
- except ImportError:
- raise RuntimeError(
- "Parakeet C++ bindings are not importable. Build `_parakeet_tdt_cpp` "
- "with CMake and add the build output directory to PYTHONPATH."
- ) from error
-
- self.transcriber = Transcriber(
- str(model_dir),
- provider,
- str(ep_cache),
- str(ep_context_dir),
- "auto",
- )
-
- def transcribe_audio(self, audio: Audio) -> dict[str, Any]:
- samples = np.ascontiguousarray(audio.samples, dtype=np.float32)
- return dict(self.transcriber.transcribe(samples, int(audio.sampling_rate)))
-
- def transcribe_file(self, path: str | Path) -> dict[str, Any]:
- return dict(self.transcriber.transcribe_file(str(path)))
diff --git a/asr/rnnt/python/rnnt/parakeet_tdt/nemo_backend.py b/asr/rnnt/python/rnnt/parakeet_tdt/nemo_backend.py
deleted file mode 100644
index 8b4c3e8..0000000
--- a/asr/rnnt/python/rnnt/parakeet_tdt/nemo_backend.py
+++ /dev/null
@@ -1,107 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-import json
-import os
-import tempfile
-from pathlib import Path
-from typing import Any
-
-import soundfile as sf
-import torch
-from torch import nn
-
-from rnnt.audio import Audio
-
-artifacts_base = os.environ.get("DIN_DEPLOY_ARTIFACTS_PATH")
-ARTIFACTS_DIR = Path(artifacts_base) / "parakeet" if artifacts_base else Path("artifacts") / "parakeet"
-MODEL_DIR = ARTIFACTS_DIR / "model"
-ONNX_DIR = ARTIFACTS_DIR / "onnx"
-DEFAULT_MODEL_ID = "nvidia/parakeet-tdt-0.6b-v3"
-DEFAULT_MODEL_DIR = MODEL_DIR
-
-
-def load_nemo_model(
- model_id: str = DEFAULT_MODEL_ID,
- *,
- device: str | torch.device | None = None,
- dtype: torch.dtype = torch.float16,
-) -> nn.Module:
- target_device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
- model = _load_nemo_asr_model(model_id, target_device)
- model = model.to(device=target_device, dtype=dtype)
- model.eval()
- return model
-
-
-def transcribe_paths(model: nn.Module, paths: list[str | Path], batch_size: int = 1) -> list[str]:
- results = model.transcribe([str(path) for path in paths], batch_size=batch_size)
- return [_result_text(result) for result in results]
-
-
-def transcribe_audio(model: nn.Module, audio: Audio) -> str:
- with tempfile.NamedTemporaryFile(suffix=".wav") as handle:
- sf.write(handle.name, audio.samples, audio.sampling_rate)
- return transcribe_paths(model, [handle.name])[0]
-
-
-def apply_attention(model: nn.Module, attention: str, att_context_size: list[int] | None) -> None:
- if not hasattr(model, "change_attention_model"):
- return
- if attention == "rel_pos":
- model.change_attention_model("rel_pos", None)
- else:
- model.change_attention_model("rel_pos_local_attn", att_context_size or [256, 256])
-
-
-def configure_encoder_sdpa(encoder: nn.Module) -> None:
- for module in [encoder, *getattr(encoder, "layers", [])]:
- if hasattr(module, "use_pytorch_sdpa"):
- module.use_pytorch_sdpa = True
- self_attn = getattr(module, "self_attn", None)
- if self_attn is not None and hasattr(self_attn, "use_pytorch_sdpa"):
- self_attn.use_pytorch_sdpa = True
-
-
-def write_tokenizer_artifacts(model: nn.Module, output: Path) -> None:
- tokenizer = getattr(model, "tokenizer", None)
- if tokenizer is None:
- raise RuntimeError("NeMo Parakeet model does not expose a tokenizer.")
-
- sentencepiece = getattr(tokenizer, "tokenizer", tokenizer)
- if hasattr(sentencepiece, "get_piece_size") and hasattr(sentencepiece, "id_to_piece"):
- vocab = {sentencepiece.id_to_piece(index): index for index in range(sentencepiece.get_piece_size())}
- (output / "tokenizer.json").write_text(json.dumps({"model": {"vocab": vocab}}, indent=2), encoding="utf-8")
- return
- if hasattr(tokenizer, "tokenizer") and hasattr(tokenizer.tokenizer, "save"):
- tokenizer.tokenizer.save(str(output / "tokenizer.json"))
- return
- if hasattr(tokenizer, "save_to"):
- tokenizer.save_to(str(output / "tokenizer.json"))
- return
- if hasattr(tokenizer, "vocab_path"):
- vocab_path = Path(tokenizer.vocab_path)
- if vocab_path.is_file():
- (output / "tokenizer.json").write_bytes(vocab_path.read_bytes())
- return
- raise RuntimeError("Unable to write tokenizer.json from NeMo Parakeet tokenizer.")
-
-
-def _load_nemo_asr_model(model_id: str, device: torch.device) -> nn.Module:
- try:
- from nemo.collections import asr as nemo_asr
- except ImportError as exc:
- raise RuntimeError("Parakeet requires the optional `nemo_toolkit[asr]` dependency.") from exc
-
- model_path = Path(model_id)
- if model_path.is_file() and model_path.suffix == ".nemo":
- return nemo_asr.models.ASRModel.restore_from(str(model_path), map_location=device)
- return nemo_asr.models.ASRModel.from_pretrained(model_name=model_id, map_location=device)
-
-
-def _result_text(result: Any) -> str:
- if hasattr(result, "text"):
- return str(result.text).strip()
- return str(result).strip()
diff --git a/asr/rnnt/python/rnnt/parakeet_tdt/timestamps.py b/asr/rnnt/python/rnnt/parakeet_tdt/timestamps.py
deleted file mode 100644
index cde31f3..0000000
--- a/asr/rnnt/python/rnnt/parakeet_tdt/timestamps.py
+++ /dev/null
@@ -1,168 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-import re
-from collections.abc import Iterable
-
-Timestamp = dict[str, float | int | str | list[str]]
-
-
-def build_timestamps(text: str, token_timestamps: Iterable[dict]) -> dict[str, list[Timestamp]]:
- char = _chars_from_tokens(text, token_timestamps)
- word = _words_from_chars(text, char)
- segment = _segments_from_words(text, word)
- return {"char": char, "word": word, "segment": segment}
-
-
-def offset_timestamps(timestamp: dict[str, list[Timestamp]], seconds: float) -> dict[str, list[Timestamp]]:
- if seconds == 0:
- return timestamp
-
- shifted: dict[str, list[Timestamp]] = {}
- for level, stamps in timestamp.items():
- shifted[level] = []
- for stamp in stamps:
- next_stamp = dict(stamp)
- start = next_stamp["start"]
- end = next_stamp["end"]
- assert not isinstance(start, list)
- assert not isinstance(end, list)
- next_stamp["start"] = round(float(start) + seconds, 3)
- next_stamp["end"] = round(float(end) + seconds, 3)
- shifted[level].append(next_stamp)
- return shifted
-
-
-def merge_timestamps(
- parts: list[dict[str, list[Timestamp]]],
-) -> dict[str, list[Timestamp]]:
- merged = {"char": [], "word": [], "segment": []}
- for part in parts:
- for level in merged:
- merged[level].extend(part.get(level, []))
- return merged
-
-
-def _chars_from_tokens(text: str, token_timestamps: Iterable[dict]) -> list[Timestamp]:
- stamps: list[Timestamp] = []
-
- for token_stamp in token_timestamps:
- token = str(token_stamp.get("token", ""))
- clean_token = _clean_token(token)
- if not token:
- continue
-
- token_text = clean_token if clean_token else token
- if not token_text:
- continue
-
- start_offset = int(token_stamp.get("start_offset", 0))
- end_offset = int(token_stamp.get("end_offset", start_offset))
- stamps.append(
- {
- "char": [token_text],
- "token": token,
- "token_id": int(token_stamp["token_id"]) if "token_id" in token_stamp else -1,
- "start_offset": start_offset,
- "end_offset": end_offset,
- "start": round(float(token_stamp["start"]), 3),
- "end": round(float(token_stamp["end"]), 3),
- }
- )
-
- return stamps
-
-
-def _words_from_chars(text: str, char_timestamps: list[Timestamp]) -> list[Timestamp]:
- words: list[Timestamp] = []
- text_positions = _token_timestamps_by_text_position(text, char_timestamps)
-
- for match in re.finditer(r"\S+", text):
- covered = [stamp for position, stamp in text_positions if match.start() <= position < match.end()]
- if not covered:
- continue
- words.append(
- {
- "word": match.group(0),
- "start_offset": covered[0]["start_offset"],
- "end_offset": covered[-1]["end_offset"],
- "start": covered[0]["start"],
- "end": covered[-1]["end"],
- }
- )
-
- return words
-
-
-def _segments_from_words(text: str, word_timestamps: list[Timestamp]) -> list[Timestamp]:
- if not word_timestamps:
- return []
-
- segments: list[Timestamp] = []
- current_words: list[Timestamp] = []
-
- for stamp in word_timestamps:
- current_words.append(stamp)
- if str(stamp["word"]).endswith((".", "!", "?")):
- segments.append(_segment(current_words))
- current_words = []
-
- if current_words:
- segments.append(_segment(current_words))
-
- if len(segments) == 1:
- segments[0]["segment"] = text
-
- return segments
-
-
-def _segment(words: list[Timestamp]) -> Timestamp:
- return {
- "segment": " ".join(str(word["word"]) for word in words),
- "start_offset": words[0]["start_offset"],
- "end_offset": words[-1]["end_offset"],
- "start": words[0]["start"],
- "end": words[-1]["end"],
- }
-
-
-def _clean_token(token: str) -> str:
- return token.replace("▁", " ").replace("", "")
-
-
-def _advance_to_char(text: str, start: int, expected: str) -> int:
- expected = expected.casefold()
- for index in range(start, len(text)):
- if text[index].isspace():
- continue
- if text[index].casefold() == expected:
- return index
- return len(text)
-
-
-def _token_timestamps_by_text_position(
- text: str,
- char_timestamps: list[Timestamp],
-) -> list[tuple[int, Timestamp]]:
- by_position: list[tuple[int, Timestamp]] = []
- cursor = 0
- for stamp in char_timestamps:
- token_chars = _token_chars(stamp)
- token_start: int | None = None
- for char in token_chars:
- cursor = _advance_to_char(text, cursor, char)
- if cursor >= len(text):
- break
- token_start = cursor if token_start is None else token_start
- cursor += 1
- if token_start is not None:
- by_position.append((token_start, stamp))
- return by_position
-
-
-def _token_chars(stamp: Timestamp) -> list[str]:
- raw_char = stamp["char"]
- text = "".join(raw_char) if isinstance(raw_char, list) else str(raw_char)
- return [char for char in text if not char.isspace()]
diff --git a/asr/rnnt/python/rnnt/validation/__init__.py b/asr/rnnt/python/rnnt/validation/__init__.py
deleted file mode 100644
index 4025ffb..0000000
--- a/asr/rnnt/python/rnnt/validation/__init__.py
+++ /dev/null
@@ -1,4 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-"""Small validation helpers for the demo repo."""
diff --git a/asr/rnnt/python/rnnt/validation/__main__.py b/asr/rnnt/python/rnnt/validation/__main__.py
deleted file mode 100644
index c69d090..0000000
--- a/asr/rnnt/python/rnnt/validation/__main__.py
+++ /dev/null
@@ -1,428 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-import argparse
-import json
-import random
-import time
-from collections.abc import Iterable, Iterator
-from pathlib import Path
-from types import SimpleNamespace
-from typing import Any, TypeVar
-
-import numpy as np
-
-from .data import (
- DEFAULT_MAX_SAMPLES_PER_LANGUAGE,
- DEFAULT_SPLIT,
- ValidationItem,
- iter_audio_dir,
- iter_fleurs,
- iter_librispeech_long,
- iter_manifest,
- maybe_group_long_form,
- parse_languages,
- parse_max_samples,
-)
-from .metrics import ValidationSummary, normalize_text, student_losses
-
-T = TypeVar("T")
-FAST_SAMPLES_PER_LANGUAGE = 20
-
-
-def main() -> None:
- args = parse_args()
- run(args)
-
-
-def run(args: argparse.Namespace) -> None:
- output = args.output
- output.mkdir(parents=True, exist_ok=True)
- set_seed(args.seed)
- validate_args(args)
-
- cpp_backend = None
- model = None
- tokenizer = None
- if args.backend == "cpp":
- if args.quantized_modelopt_path is not None:
- raise ValueError("--quantized-modelopt-path is only supported with --backend python.")
- if args.student_loss:
- raise ValueError("--student-loss is only supported with --backend python.")
- if args.max_feature_frames is not None:
- raise ValueError("--max-feature-frames is only supported with --backend python.")
- cpp_backend = load_cpp_backend(args)
- else:
- model, tokenizer = load_python_model(args)
-
- teacher = None
- if args.student_loss:
- if args.asr_model != "parakeet":
- raise ValueError("--student-loss is only supported with --asr-model parakeet.")
- import torch
-
- from rnnt.parakeet_tdt.nemo_backend import DEFAULT_MODEL_ID, load_nemo_model
-
- teacher = load_nemo_model(args.model or DEFAULT_MODEL_ID, device=args.device, dtype=torch.float32)
- apply_attention(teacher, args.attention, args.att_context_size)
-
- summary = ValidationSummary()
- prediction_path = output / "predictions.jsonl"
- with prediction_path.open("w", encoding="utf-8") as predictions:
- items = progress(limited_items(load_items(args), args.limit), total=progress_total(args), desc="validating")
- for item in items:
- started = time.perf_counter()
- result = transcribe_item(model, tokenizer, cpp_backend, item, args)
- elapsed = time.perf_counter() - started
-
- has_reference = bool(normalize_text(item.text))
- if has_reference:
- wer = summary.add(
- language=item.language,
- duration=item.duration,
- transcribe_seconds=elapsed,
- reference=item.text,
- hypothesis=result.text,
- )
- else:
- summary.add_unscored(duration=item.duration, transcribe_seconds=elapsed)
- wer = None
- record: dict[str, Any] = {
- "id": item.id,
- "language": item.language,
- "duration": item.duration,
- "reference_available": has_reference,
- "reference": item.text,
- "hypothesis": result.text,
- "reference_normalized": normalize_text(item.text),
- "hypothesis_normalized": normalize_text(result.text),
- "wer": wer.wer if wer else None,
- "errors": wer.errors if wer else None,
- "reference_words": wer.reference_words if wer else 0,
- "transcribe_seconds": elapsed,
- "rtf": elapsed / item.duration if item.duration else 0.0,
- }
- if teacher is not None:
- record["student_loss"] = student_losses(teacher, model, item.audio)
- predictions.write(json.dumps(record, ensure_ascii=False) + "\n")
- predictions.flush()
-
- write_summary(output / "summary.json", args, summary)
- print(json.dumps(summary.as_dict(), indent=2))
-
-
-def load_items(args: argparse.Namespace) -> Iterator[ValidationItem]:
- if args.manifest is not None:
- items: Iterable[ValidationItem] = iter_manifest(args.manifest)
- elif args.audio_dir is not None:
- items = iter_audio_dir(args.audio_dir)
- elif args.dataset == "fleurs":
- items = iter_fleurs(
- parse_languages(args.languages),
- split=args.split,
- max_samples_per_language=effective_max_samples_per_language(args),
- )
- else:
- items = iter_librispeech_long(
- split=args.split,
- max_samples=args.limit,
- )
- yield from maybe_group_long_form(items, args.long_form_target_seconds)
-
-
-def limited_items(items: Iterable[ValidationItem], limit: int | None) -> Iterator[ValidationItem]:
- for index, item in enumerate(items):
- if limit is not None and index >= limit:
- break
- yield item
-
-
-def progress[T](items: Iterable[T], *, total: int | None, desc: str) -> Iterable[T]:
- try:
- from tqdm import tqdm
- except ImportError:
- return items
- return tqdm(items, total=total, desc=desc, unit="utt")
-
-
-def progress_total(args: argparse.Namespace) -> int | None:
- if args.limit is not None:
- return args.limit
- if args.manifest is not None or args.audio_dir is not None or args.long_form_target_seconds:
- return None
- if args.dataset == "librispeech-long":
- return args.limit
- max_samples = effective_max_samples_per_language(args)
- if max_samples is None:
- return None
- return len(parse_languages(args.languages)) * max_samples
-
-
-def load_python_model(args: argparse.Namespace):
- dtype = parse_dtype(args.dtype)
- if args.asr_model == "parakeet":
- from rnnt.parakeet_tdt.nemo_backend import DEFAULT_MODEL_ID, load_nemo_model
-
- model = load_nemo_model(
- args.model or DEFAULT_MODEL_ID,
- device=args.device,
- dtype=dtype,
- )
- tokenizer = None
- apply_attention(model, args.attention, args.att_context_size)
- else:
- from rnnt.nemotron_asr.nemo_backend import DEFAULT_ATT_CONTEXT_SIZE, DEFAULT_MODEL_ID, load_nemo_model
-
- model = load_nemo_model(
- args.model or DEFAULT_MODEL_ID,
- device=args.device,
- target_lang=args.target_lang,
- att_context_size=DEFAULT_ATT_CONTEXT_SIZE,
- dtype=dtype,
- )
- tokenizer = None
-
- if args.quantized_modelopt_path is not None:
- from .quantization import restore_modelopt_quantized_model
-
- model, _ = restore_modelopt_quantized_model(
- model,
- args.quantized_modelopt_path,
- summary_path=args.output / "quant_summary.txt",
- )
- return model, tokenizer
-
-
-def load_cpp_backend(args: argparse.Namespace):
- if args.asr_model == "parakeet":
- from rnnt.parakeet_tdt.cpp_backend import CppBackend
-
- return CppBackend(
- model_dir=args.cpp_model_dir,
- provider=args.cpp_provider,
- ep_cache=args.cpp_ep_cache,
- ep_context_dir=args.cpp_ep_context_dir,
- )
-
- from rnnt.nemotron_asr.cpp_backend import CppBackend
-
- return CppBackend(
- model_dir=args.cpp_model_dir,
- provider=args.cpp_provider,
- ep_cache=args.cpp_ep_cache,
- ep_context_dir=args.cpp_ep_context_dir,
- lang_id=args.target_lang,
- )
-
-
-def transcribe_item(model, tokenizer, cpp_backend, item: ValidationItem, args: argparse.Namespace):
- if cpp_backend is not None:
- raw = cpp_backend.transcribe_audio(item.audio)
- return SimpleNamespace(text=str(raw["text"]), raw=raw, timestamp={})
- assert model is not None
- if args.asr_model == "nemotron":
- return transcribe_nemotron_item(model, tokenizer, item, args)
-
- from rnnt.parakeet_tdt.nemo_backend import transcribe_audio
-
- del tokenizer
- return SimpleNamespace(text=transcribe_audio(model, item.audio), timestamp={})
-
-
-def transcribe_nemotron_item(model, tokenizer, item: ValidationItem, args: argparse.Namespace):
- from rnnt.nemotron_asr.nemo_backend import strip_lang_tags, transcribe_audio
-
- del tokenizer
- text = strip_lang_tags(transcribe_audio(model, item.audio, args.target_lang))
- return SimpleNamespace(text=text.strip(), timestamp={})
-
-
-def apply_attention(model, attention: str, att_context_size: str | None) -> None:
- from rnnt.parakeet_tdt.nemo_backend import apply_attention as apply_parakeet_attention
-
- context = None if attention == "rel_pos" else parse_att_context_size(att_context_size)
- apply_parakeet_attention(model, attention, context)
-
-
-def parse_att_context_size(value: str | None) -> list[int]:
- if not value:
- return [256, 256]
- parts = [int(part.strip()) for part in value.split(",") if part.strip()]
- if len(parts) != 2:
- raise ValueError("--att-context-size must have two comma-separated integers.")
- return parts
-
-
-def parse_dtype(value: str):
- import torch
-
- if value == "float32":
- return torch.float32
- if value == "float16":
- return torch.float16
- raise ValueError(f"Unsupported dtype: {value}")
-
-
-def validate_args(args: argparse.Namespace) -> None:
- if args.asr_model == "nemotron":
- if args.student_loss:
- raise ValueError("--student-loss is only supported with --asr-model parakeet.")
- if args.max_feature_frames is not None:
- raise ValueError("--max-feature-frames is only supported with --asr-model parakeet.")
- if args.window_seconds is not None:
- raise ValueError("--window-seconds is only supported with --asr-model parakeet.")
- elif args.backend == "python":
- if args.max_feature_frames is not None:
- raise ValueError("--max-feature-frames belonged to the removed custom Parakeet sampler.")
- if args.window_seconds is not None:
- raise ValueError("--window-seconds belonged to the removed custom Parakeet sampler.")
-
-
-def effective_max_samples_per_language(args: argparse.Namespace) -> int | None:
- if args.fast:
- return args.fast_samples_per_language
- return parse_max_samples(args.max_samples_per_language)
-
-
-def set_seed(seed: int) -> None:
- random.seed(seed)
- np.random.seed(seed)
- try:
- import torch
- except ImportError:
- return
- torch.manual_seed(seed)
- if torch.cuda.is_available():
- torch.cuda.manual_seed_all(seed)
- torch.use_deterministic_algorithms(True, warn_only=True)
-
-
-def write_summary(
- path: Path,
- args: argparse.Namespace,
- summary: ValidationSummary,
-) -> None:
- config = {
- "dataset": args.dataset,
- "backend": args.backend,
- "manifest": str(args.manifest) if args.manifest else None,
- "audio_dir": str(args.audio_dir) if args.audio_dir else None,
- "languages": (
- parse_languages(args.languages)
- if args.manifest is None and args.audio_dir is None and args.dataset == "fleurs"
- else None
- ),
- "split": args.split if args.manifest is None and args.audio_dir is None else None,
- "max_samples_per_language": (
- args.max_samples_per_language
- if args.manifest is None and args.audio_dir is None and args.dataset == "fleurs"
- else None
- ),
- "fast": args.fast,
- "fast_samples_per_language": args.fast_samples_per_language if args.fast else None,
- "effective_max_samples_per_language": (
- effective_max_samples_per_language(args)
- if args.manifest is None and args.audio_dir is None and args.dataset == "fleurs"
- else None
- ),
- "attention": args.attention,
- "att_context_size": args.att_context_size,
- "window_seconds": args.window_seconds,
- "max_feature_frames": args.max_feature_frames,
- "budget_left_context_seconds": args.budget_left_context_seconds,
- "budget_right_context_seconds": args.budget_right_context_seconds,
- "budget_chunk_seconds": args.budget_chunk_seconds,
- "long_form_target_seconds": args.long_form_target_seconds,
- "student_loss": args.student_loss,
- "asr_model": args.asr_model,
- "model": args.model,
- "model_dir": str(args.model_dir) if args.model_dir else None,
- "device": args.device,
- "dtype": args.dtype,
- "seed": args.seed,
- "quantized_modelopt_path": str(args.quantized_modelopt_path) if args.quantized_modelopt_path else None,
- "target_lang": args.target_lang if args.asr_model == "nemotron" else None,
- "cpp_model_dir": str(args.cpp_model_dir) if args.backend == "cpp" else None,
- "cpp_provider": args.cpp_provider if args.backend == "cpp" else None,
- "cpp_ep_cache": str(args.cpp_ep_cache) if args.backend == "cpp" else None,
- "cpp_ep_context_dir": str(args.cpp_ep_context_dir) if args.backend == "cpp" else None,
- }
- payload = {"config": config, "summary": summary.as_dict()}
- path.write_text(
- json.dumps(payload, indent=2),
- encoding="utf-8",
- )
-
-
-def parse_args() -> argparse.Namespace:
- parser = argparse.ArgumentParser(description="Run a small RNNT validation pass.")
- source = parser.add_mutually_exclusive_group()
- source.add_argument("--dataset", choices=("fleurs", "librispeech-long"), default="fleurs")
- source.add_argument("--manifest", type=Path)
- source.add_argument("--audio-dir", type=Path)
- parser.add_argument("--languages", default=None, help="Comma-separated FLEURS configs.")
- parser.add_argument("--split", default=DEFAULT_SPLIT)
- parser.add_argument(
- "--max-samples-per-language",
- default=str(DEFAULT_MAX_SAMPLES_PER_LANGUAGE),
- help="FLEURS cap per language; use 0 or none for the full split.",
- )
- parser.add_argument("--limit", type=int, default=None, help="Optional global item limit.")
- parser.add_argument(
- "--fast",
- action=argparse.BooleanOptionalAction,
- default=False,
- help="Shortcut for a small validation run over each configured language.",
- )
- parser.add_argument(
- "--fast-samples-per-language",
- type=int,
- default=FAST_SAMPLES_PER_LANGUAGE,
- )
- parser.add_argument("--output", type=Path, required=True)
- parser.add_argument("--backend", choices=("python", "cpp"), default="python")
- parser.add_argument("--asr-model", choices=("parakeet", "nemotron"), default="parakeet")
- parser.add_argument("--model", default=None, help="Optional Hugging Face model id or local model path.")
- parser.add_argument("--model-dir", type=Path, default=None)
- parser.add_argument("--quantized-modelopt-path", type=Path, default=None)
- parser.add_argument("--device", default="cuda")
- parser.add_argument("--dtype", choices=("float16", "float32"), default="float16")
- parser.add_argument("--seed", type=int, default=1234)
- parser.add_argument("--attention", choices=("rel_pos", "rel_pos_local_attn"), default="rel_pos")
- parser.add_argument("--att-context-size", default=None)
- parser.add_argument("--window-seconds", type=float, default=None)
- parser.add_argument("--max-feature-frames", type=int, default=None)
- parser.add_argument("--budget-left-context-seconds", type=float, default=10.0)
- parser.add_argument("--budget-right-context-seconds", type=float, default=5.0)
- parser.add_argument(
- "--budget-chunk-seconds",
- type=float,
- default=None,
- help="Optional core chunk override. By default, use the largest chunk that fits the frame budget.",
- )
- parser.add_argument("--long-form-target-seconds", type=float, default=None)
- parser.add_argument("--student-loss", action="store_true")
- parser.add_argument("--target-lang", default="auto")
- parser.add_argument("--cpp-model-dir", type=Path, default=None)
- parser.add_argument("--cpp-provider", choices=("trt-rtx", "cpu"), default="trt-rtx")
- parser.add_argument("--cpp-ep-cache", type=Path, default=None)
- parser.add_argument("--cpp-ep-context-dir", type=Path, default=None)
-
- args = parser.parse_args()
- resolve_defaults(args)
- return args
-
-
-def resolve_defaults(args: argparse.Namespace) -> None:
- if args.cpp_model_dir is None:
- args.cpp_model_dir = Path("artifacts") / args.asr_model / "onnx"
- if args.cpp_ep_cache is None:
- args.cpp_ep_cache = Path("artifacts") / args.asr_model / "trt_rtx_cache"
- if args.cpp_ep_context_dir is None:
- args.cpp_ep_context_dir = Path("artifacts") / args.asr_model / "ep_context"
-
-
-if __name__ == "__main__":
- main()
diff --git a/asr/rnnt/python/rnnt/validation/data.py b/asr/rnnt/python/rnnt/validation/data.py
deleted file mode 100644
index 9240fc7..0000000
--- a/asr/rnnt/python/rnnt/validation/data.py
+++ /dev/null
@@ -1,259 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-import json
-from collections.abc import Iterable, Iterator
-from dataclasses import dataclass
-from difflib import get_close_matches
-from io import BytesIO
-from pathlib import Path
-from typing import Any
-
-import numpy as np
-import soundfile as sf
-
-from rnnt.audio import Audio, load_audio, resample
-
-DEFAULT_FLEURS_LANGUAGES = ("en_us", "de_de", "fr_fr", "es_419", "it_it")
-LIBRISPEECH_LONG_DATASET = "distil-whisper/librispeech_long"
-LIBRISPEECH_LONG_CONFIG = "clean"
-DEFAULT_SPLIT = "validation"
-DEFAULT_MAX_SAMPLES_PER_LANGUAGE = 50
-AUDIO_EXTENSIONS = {".flac", ".m4a", ".mp3", ".ogg", ".wav"}
-
-
-@dataclass(frozen=True)
-class ValidationItem:
- id: str
- audio: Audio
- text: str
- language: str = ""
-
- @property
- def duration(self) -> float:
- return self.audio.duration
-
-
-def parse_languages(value: str | None) -> list[str]:
- if not value:
- return list(DEFAULT_FLEURS_LANGUAGES)
- return [part.strip() for part in value.split(",") if part.strip()]
-
-
-def parse_max_samples(value: str) -> int | None:
- normalized = value.strip().lower()
- if normalized in {"none", "all"}:
- return None
- count = int(normalized)
- return None if count <= 0 else count
-
-
-def iter_manifest(path: Path, target_sampling_rate: int = 16_000) -> Iterator[ValidationItem]:
- root = path.parent
- with path.open("r", encoding="utf-8") as handle:
- for line_number, line in enumerate(handle, start=1):
- if not line.strip():
- continue
- record = json.loads(line)
- audio_path = Path(record["audio"])
- if not audio_path.is_absolute():
- audio_path = root / audio_path
- yield ValidationItem(
- id=str(record.get("id", f"{path.stem}-{line_number}")),
- audio=load_audio(audio_path, target_sampling_rate),
- text=str(record.get("text", "")),
- language=str(record.get("language", "")),
- )
-
-
-def iter_audio_dir(path: Path, target_sampling_rate: int = 16_000) -> Iterator[ValidationItem]:
- for audio_path in sorted(path.rglob("*")):
- if not audio_path.is_file() or audio_path.suffix.lower() not in AUDIO_EXTENSIONS:
- continue
- relative = audio_path.relative_to(path)
- language = relative.parts[0] if len(relative.parts) > 1 else ""
- text = read_audio_dir_groundtruth(path, audio_path)
- yield ValidationItem(
- id=relative.as_posix(),
- audio=load_audio(audio_path, target_sampling_rate),
- text=text,
- language=language,
- )
-
-
-def read_audio_dir_groundtruth(root: Path, audio_path: Path) -> str:
- relative = audio_path.relative_to(root)
- candidates = [audio_path.with_suffix(".txt")]
- groundtruth_dir = None
- if len(relative.parts) > 1:
- language_root = root / relative.parts[0]
- groundtruth_dir = language_root / "groundtruth"
- candidates.append(groundtruth_dir / f"{audio_path.stem}.txt")
-
- for candidate in candidates:
- if candidate.is_file():
- return candidate.read_text(encoding="utf-8").strip()
- if groundtruth_dir is not None and groundtruth_dir.is_dir():
- stems = {candidate.stem: candidate for candidate in groundtruth_dir.glob("*.txt")}
- matches = get_close_matches(audio_path.stem, stems.keys(), n=1, cutoff=0.85)
- if matches:
- return stems[matches[0]].read_text(encoding="utf-8").strip()
- return ""
-
-
-def iter_fleurs(
- languages: Iterable[str],
- *,
- split: str = DEFAULT_SPLIT,
- max_samples_per_language: int | None = DEFAULT_MAX_SAMPLES_PER_LANGUAGE,
- target_sampling_rate: int = 16_000,
-) -> Iterator[ValidationItem]:
- try:
- from datasets import Audio as HfAudio
- from datasets import load_dataset
- except ImportError as error:
- raise RuntimeError("FLEURS validation requires `pip install -e .[validation]`.") from error
-
- for language in languages:
- dataset = load_dataset("google/fleurs", language, split=split, streaming=True)
- dataset = dataset.cast_column("audio", HfAudio(decode=False))
- for index, record in enumerate(dataset):
- if max_samples_per_language is not None and index >= max_samples_per_language:
- break
- yield item_from_hf_record(
- record,
- language=language,
- fallback_id=f"{language}-{split}-{index}",
- target_sampling_rate=target_sampling_rate,
- )
-
-
-def iter_librispeech_long(
- *,
- split: str = DEFAULT_SPLIT,
- max_samples: int | None = None,
- target_sampling_rate: int = 16_000,
-) -> Iterator[ValidationItem]:
- try:
- from datasets import Audio as HfAudio
- from datasets import load_dataset
- except ImportError as error:
- raise RuntimeError("LibriSpeech long validation requires `pip install -e .[validation]`.") from error
-
- dataset = load_dataset(
- LIBRISPEECH_LONG_DATASET,
- LIBRISPEECH_LONG_CONFIG,
- split=split,
- streaming=True,
- )
- dataset = dataset.cast_column("audio", HfAudio(decode=False))
- for index, record in enumerate(dataset):
- if max_samples is not None and index >= max_samples:
- break
- yield item_from_hf_record(
- record,
- language="en",
- fallback_id=f"librispeech-long-{split}-{index}",
- target_sampling_rate=target_sampling_rate,
- )
-
-
-def item_from_hf_record(
- record: dict[str, Any],
- *,
- language: str,
- fallback_id: str,
- target_sampling_rate: int = 16_000,
-) -> ValidationItem:
- audio_record = record["audio"]
- audio = audio_from_hf_record(audio_record, target_sampling_rate)
-
- text = record.get("transcription", record.get("raw_transcription", ""))
- return ValidationItem(
- id=str(record.get("id", fallback_id)),
- audio=audio,
- text=str(text),
- language=language,
- )
-
-
-def audio_from_hf_record(audio_record: dict[str, Any], target_sampling_rate: int) -> Audio:
- if "array" in audio_record and audio_record["array"] is not None:
- samples = np.asarray(audio_record["array"], dtype=np.float32)
- sampling_rate = int(audio_record["sampling_rate"])
- if sampling_rate != target_sampling_rate:
- samples = resample(samples, sampling_rate, target_sampling_rate)
- sampling_rate = target_sampling_rate
- return Audio(samples=samples, sampling_rate=sampling_rate)
-
- if audio_record.get("path") and Path(audio_record["path"]).is_file():
- return load_audio(audio_record["path"], target_sampling_rate)
-
- if audio_record.get("bytes"):
- samples, sampling_rate = sf.read(BytesIO(audio_record["bytes"]), dtype="float32", always_2d=False)
- if samples.ndim == 2:
- samples = samples.mean(axis=1)
- if sampling_rate != target_sampling_rate:
- samples = resample(samples, sampling_rate, target_sampling_rate)
- sampling_rate = target_sampling_rate
- return Audio(np.asarray(samples, dtype=np.float32), sampling_rate)
-
- raise ValueError("HF audio record must contain array, path, or bytes.")
-
-
-def maybe_group_long_form(
- items: Iterable[ValidationItem],
- target_seconds: float | None,
- *,
- silence_seconds: float = 0.25,
-) -> Iterator[ValidationItem]:
- if target_seconds is None or target_seconds <= 0:
- yield from items
- return
-
- pending: list[ValidationItem] = []
- pending_seconds = 0.0
- pending_language: str | None = None
-
- for item in items:
- if pending and pending_language != item.language:
- yield combine_items(pending, silence_seconds)
- pending = []
- pending_seconds = 0.0
-
- pending.append(item)
- pending_language = item.language
- pending_seconds += item.duration
-
- if pending_seconds >= target_seconds:
- yield combine_items(pending, silence_seconds)
- pending = []
- pending_seconds = 0.0
- pending_language = None
-
- if pending:
- yield combine_items(pending, silence_seconds)
-
-
-def combine_items(items: list[ValidationItem], silence_seconds: float = 0.25) -> ValidationItem:
- if not items:
- raise ValueError("Cannot combine an empty item list.")
-
- sampling_rate = items[0].audio.sampling_rate
- silence = np.zeros(max(0, round(silence_seconds * sampling_rate)), dtype=np.float32)
- pieces: list[np.ndarray] = []
- for index, item in enumerate(items):
- if item.audio.sampling_rate != sampling_rate:
- raise ValueError("Long-form grouping requires matching sampling rates.")
- if index:
- pieces.append(silence)
- pieces.append(item.audio.samples)
-
- return ValidationItem(
- id="+".join(item.id for item in items),
- audio=Audio(np.concatenate(pieces).astype(np.float32, copy=False), sampling_rate),
- text=" ".join(item.text for item in items),
- language=items[0].language,
- )
diff --git a/asr/rnnt/python/rnnt/validation/metrics.py b/asr/rnnt/python/rnnt/validation/metrics.py
deleted file mode 100644
index e41a4fd..0000000
--- a/asr/rnnt/python/rnnt/validation/metrics.py
+++ /dev/null
@@ -1,305 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-import unicodedata
-from dataclasses import dataclass, field
-
-from rnnt.audio import Audio
-
-
-def normalize_text(text: str) -> str:
- chars = []
- for char in text.casefold():
- if unicodedata.category(char).startswith("P"):
- chars.append(" ")
- else:
- chars.append(char)
- return " ".join("".join(chars).split())
-
-
-@dataclass
-class WerResult:
- wer: float
- errors: int
- reference_words: int
-
-
-@dataclass
-class WerAccumulator:
- errors: int = 0
- reference_words: int = 0
- count: int = 0
-
- def add(self, reference: str, hypothesis: str) -> WerResult:
- result = compute_wer(reference, hypothesis)
- self.errors += result.errors
- self.reference_words += result.reference_words
- self.count += 1
- return result
-
- @property
- def wer(self) -> float:
- if self.reference_words == 0:
- return 0.0 if self.errors == 0 else 1.0
- return self.errors / self.reference_words
-
- def as_dict(self) -> dict[str, float | int]:
- return {
- "wer": self.wer,
- "errors": self.errors,
- "reference_words": self.reference_words,
- "count": self.count,
- }
-
-
-@dataclass
-class ValidationSummary:
- total: WerAccumulator = field(default_factory=WerAccumulator)
- by_language: dict[str, WerAccumulator] = field(default_factory=dict)
- audio_seconds: float = 0.0
- transcribe_seconds: float = 0.0
- unscored_count: int = 0
-
- def add(
- self,
- *,
- language: str,
- duration: float,
- transcribe_seconds: float,
- reference: str,
- hypothesis: str,
- ) -> WerResult:
- result = self.total.add(reference, hypothesis)
- self.by_language.setdefault(language or "unknown", WerAccumulator()).add(reference, hypothesis)
- self.audio_seconds += duration
- self.transcribe_seconds += transcribe_seconds
- return result
-
- def add_unscored(
- self,
- *,
- duration: float,
- transcribe_seconds: float,
- ) -> None:
- self.unscored_count += 1
- self.audio_seconds += duration
- self.transcribe_seconds += transcribe_seconds
-
- def as_dict(self) -> dict:
- rtf = self.transcribe_seconds / self.audio_seconds if self.audio_seconds else 0.0
- return {
- **self.total.as_dict(),
- "unscored_count": self.unscored_count,
- "audio_seconds": self.audio_seconds,
- "transcribe_seconds": self.transcribe_seconds,
- "rtf": rtf,
- "by_language": {
- language: accumulator.as_dict() for language, accumulator in sorted(self.by_language.items())
- },
- }
-
-
-def compute_wer(reference: str, hypothesis: str) -> WerResult:
- reference_words = normalize_text(reference).split()
- hypothesis_words = normalize_text(hypothesis).split()
- errors = edit_distance(reference_words, hypothesis_words)
- total = len(reference_words)
- wer = (0.0 if errors == 0 else 1.0) if total == 0 else errors / total
- return WerResult(wer=wer, errors=errors, reference_words=total)
-
-
-def edit_distance(reference: list[str], hypothesis: list[str]) -> int:
- previous = list(range(len(hypothesis) + 1))
- for i, ref_word in enumerate(reference, start=1):
- current = [i]
- for j, hyp_word in enumerate(hypothesis, start=1):
- cost = 0 if ref_word == hyp_word else 1
- current.append(
- min(
- previous[j] + 1,
- current[j - 1] + 1,
- previous[j - 1] + cost,
- )
- )
- previous = current
- return previous[-1]
-
-
-def student_losses(
- teacher,
- student,
- audio: Audio,
- *,
- max_joint_steps: int = 256,
-) -> dict[str, float]:
- import torch
- from torch.nn import functional as F
-
- torch.set_grad_enabled(False)
- teacher_states, teacher_mask = encode_audio(teacher, audio)
- student_states, student_mask = encode_audio(student, audio)
-
- common = min(teacher_states.shape[1], student_states.shape[1])
- teacher_states = teacher_states[:, :common]
- student_states = student_states[:, :common]
- if teacher_mask is not None and student_mask is not None:
- valid_mask = teacher_mask[:, :common] & student_mask[:, :common]
- else:
- valid_mask = torch.ones(teacher_states.shape[:2], device=teacher_states.device, dtype=torch.bool)
-
- encoder_mse = F.mse_loss(
- student_states[valid_mask].float(),
- teacher_states[valid_mask].to(student_states.device).float(),
- )
- joint_kl, steps = teacher_path_joint_kl(
- teacher,
- student,
- teacher_states,
- student_states,
- int(valid_mask[0].sum().item()),
- max_steps=max_joint_steps,
- )
- return {
- "encoder_mse": float(encoder_mse.item()),
- "joint_kl": float(joint_kl),
- "joint_steps": steps,
- }
-
-
-def encode_audio(model, audio: Audio):
- import torch
-
- samples, lengths = audio_batch([audio], model.device)
- features, feature_lengths = model.preprocessor(input_signal=samples, length=lengths)
- compute_dtype = next(model.encoder.parameters()).dtype
- encoded, encoded_lengths = model.encoder(audio_signal=features.to(dtype=compute_dtype), length=feature_lengths)
- states = encoded.transpose(1, 2)
- mask = torch.arange(states.shape[1], device=states.device)[None] < encoded_lengths[:, None]
- return states, mask
-
-
-def teacher_path_joint_kl(
- teacher,
- student,
- teacher_states,
- student_states,
- valid_length: int,
- *,
- max_steps: int,
-) -> tuple[float, int]:
- import torch
- from torch.nn import functional as F
-
- teacher_state = initial_decoder_state(teacher, teacher_states.device)
- student_state = initial_decoder_state(student, student_states.device)
- time_idx = 0
- steps = 0
- kl_sum = 0.0
-
- while time_idx < valid_length and steps < max_steps:
- decoder_input = torch.tensor([[teacher_state.last_token]], device=teacher_states.device, dtype=torch.long)
- teacher_decoder, (teacher_hidden, teacher_cell) = teacher.decoder.predict(
- y=decoder_input,
- state=(teacher_state.hidden, teacher_state.cell),
- add_sos=False,
- batch_size=1,
- )
- teacher_logits = parakeet_joint_logits(teacher, teacher_states[:, time_idx : time_idx + 1], teacher_decoder)
-
- student_input = decoder_input.to(student_states.device)
- student_decoder, (student_hidden, student_cell) = student.decoder.predict(
- y=student_input,
- state=(student_state.hidden, student_state.cell),
- add_sos=False,
- batch_size=1,
- )
- student_logits = parakeet_joint_logits(student, student_states[:, time_idx : time_idx + 1], student_decoder)
-
- kl_sum += float(
- F.kl_div(
- F.log_softmax(student_logits, dim=-1),
- F.softmax(teacher_logits.to(student_logits.device), dim=-1),
- reduction="batchmean",
- ).item()
- )
-
- vocab_size = blank_token_id(teacher) + 1
- token = int(teacher_logits[..., :vocab_size].argmax(dim=-1).item())
- duration_index = int(teacher_logits[..., vocab_size:].argmax(dim=-1).item())
- skip = int(durations(teacher)[duration_index])
- steps += 1
-
- if token != blank_token_id(teacher):
- teacher_state.last_token = token
- teacher_state.hidden = teacher_hidden
- teacher_state.cell = teacher_cell
- student_state.last_token = token
- student_state.hidden = student_hidden.float()
- student_state.cell = student_cell.float()
-
- time_idx += skip if skip else 1
-
- return (kl_sum / steps if steps else 0.0), steps
-
-
-def audio_batch(audios: list[Audio], device):
- import torch
-
- max_len = max(len(audio.samples) for audio in audios)
- batch = torch.zeros(len(audios), max_len, device=device, dtype=torch.float32)
- lengths = torch.empty(len(audios), device=device, dtype=torch.long)
- for index, audio in enumerate(audios):
- samples = torch.from_numpy(audio.samples).to(device=device, dtype=torch.float32)
- batch[index, : samples.numel()] = samples
- lengths[index] = samples.numel()
- return batch, lengths
-
-
-def initial_decoder_state(model, device):
- from types import SimpleNamespace
-
- import torch
-
- blank_token = blank_token_id(model)
- encoder_dim = joint_encoder_hidden(model)
- hidden, cell = model.decoder.initialize_state(torch.zeros(1, 1, encoder_dim, device=device))
- return SimpleNamespace(last_token=blank_token, hidden=hidden, cell=cell)
-
-
-def blank_token_id(model) -> int:
- return int(model.cfg.decoder.vocab_size)
-
-
-def durations(model) -> tuple[int, ...]:
- for value in (
- getattr(getattr(model.cfg, "model_defaults", None), "tdt_durations", None),
- getattr(getattr(model.cfg, "joint", None), "durations", None),
- getattr(model.cfg, "durations", None),
- ):
- if value is not None:
- return tuple(int(item) for item in value)
- return (0, 1, 2, 3, 4)
-
-
-def joint_encoder_hidden(model) -> int:
- prednet = getattr(model.cfg.decoder, "prednet", model.cfg.decoder)
- return int(getattr(model.cfg.joint.jointnet, "encoder_hidden", prednet.pred_hidden))
-
-
-def parakeet_joint_logits(model, encoder_states, decoder_output):
- joint = model.joint
- if (
- hasattr(joint, "project_encoder")
- and hasattr(joint, "project_prednet")
- and hasattr(joint, "joint_after_projection")
- ):
- enc = joint.project_encoder(encoder_states)
- pred = joint.project_prednet(decoder_output)
- logits = joint.joint_after_projection(enc, pred)
- else:
- logits = joint(encoder_states, decoder_output)
- if logits.dim() == 4 and logits.shape[2] == 1:
- logits = logits.squeeze(2)
- return logits.float()
diff --git a/asr/rnnt/python/rnnt/validation/quantization.py b/asr/rnnt/python/rnnt/validation/quantization.py
deleted file mode 100644
index 8a3be8c..0000000
--- a/asr/rnnt/python/rnnt/validation/quantization.py
+++ /dev/null
@@ -1,466 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-import argparse
-import io
-import json
-from collections.abc import Iterable, Iterator
-from contextlib import redirect_stdout
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Any
-
-import torch
-
-from rnnt.audio import Audio
-
-from .data import DEFAULT_SPLIT, ValidationItem, iter_audio_dir, iter_fleurs, iter_manifest, parse_languages
-
-
-@dataclass(frozen=True)
-class QuantModel:
- model: torch.nn.Module
- sample_rate: int
- score_step: Any
-
-
-@dataclass(frozen=True)
-class JointLossOutput:
- token_logits: torch.Tensor
- token_targets: torch.Tensor
- duration_logits: torch.Tensor | None = None
- duration_targets: torch.Tensor | None = None
-
-
-def main() -> None:
- args = parse_args()
- run(args)
-
-
-def run(args: argparse.Namespace) -> None:
- args.output.mkdir(parents=True, exist_ok=True)
- log(f"loading {args.model} candidate model on {args.device}")
- quant_model = load_quant_model(args)
- log(f"loading calibration data from {calibration_source(args)}")
- calib_items = list(load_calibration_items(args, quant_model.sample_rate))
- if not calib_items:
- raise ValueError("No calibration audio found.")
- if len(calib_items) < args.calib_samples:
- raise ValueError(
- f"Requested {args.calib_samples} calibration samples, but only found {len(calib_items)} "
- f"from {calibration_source(args)}."
- )
- log(f"loaded {len(calib_items)} calibration samples")
-
- info = auto_quantize_encoder(
- quant_model.model,
- calib_items,
- quant_model.score_step,
- precision=args.precision,
- effective_bits=args.effective_bits,
- calib_steps=args.calib_steps,
- score_steps=args.score_steps,
- verbose=args.verbose,
- summary_path=args.output / "quant_summary.txt",
- state_path=args.output / "auto_quant_state.json",
- modelopt_path=args.output / "auto_quantize_model.pt",
- )
- (args.output / "quantization.json").write_text(json.dumps(info, indent=2), encoding="utf-8")
- print(json.dumps(info, indent=2))
-
-
-def auto_quantize_encoder(
- model,
- calib_items: list[ValidationItem],
- score_step,
- *,
- precision: str,
- effective_bits: float,
- calib_steps: int,
- score_steps: int,
- verbose: bool,
- summary_path: Path,
- state_path: Path,
- modelopt_path: Path,
-) -> dict[str, Any]:
- import modelopt.torch.opt as mto
- import modelopt.torch.quantization as mtq
-
- def forward_step(quant_model, data: ValidationItem) -> JointLossOutput:
- return score_step(quant_model, data)
-
- def loss_func(output: JointLossOutput, data: ValidationItem) -> torch.Tensor:
- loss = torch.nn.functional.cross_entropy(output.token_logits.float(), output.token_targets)
- if output.duration_logits is not None and output.duration_targets is not None:
- loss = loss + torch.nn.functional.cross_entropy(output.duration_logits.float(), output.duration_targets)
- log(f"loss: {loss.item()}")
- return loss
-
- model.eval()
- log(f"starting ModelOpt auto_quantize ({precision}, effective_bits={effective_bits})")
- model, search_state = mtq.auto_quantize(
- model,
- constraints={"effective_bits": effective_bits},
- quantization_formats=[quantization_format(mtq, precision)],
- data_loader=calib_items,
- forward_step=forward_step,
- loss_func=loss_func,
- disabled_layers=["feature_extractor*", "preprocessor*", "decoder*", "joint*"],
- num_calib_steps=min(calib_steps, len(calib_items)),
- num_score_steps=min(score_steps, len(calib_items)),
- verbose=verbose,
- )
- log(f"saving searched model to {modelopt_path}")
- mto.save(model, modelopt_path)
- log(f"writing quantization summary to {summary_path}")
- print_and_store_quant_summary(mtq, model, summary_path, verbose=verbose)
- state_path.write_text(json.dumps(jsonable(search_state), indent=2), encoding="utf-8")
- return {
- "mode": "auto-quantize-encoder",
- "precision": precision,
- "effective_bits": effective_bits,
- "calib_steps": min(calib_steps, len(calib_items)),
- "score_steps": min(score_steps, len(calib_items)),
- "sample_count": len(calib_items),
- "modelopt_model_path": str(modelopt_path),
- "auto_quant_state_path": str(state_path),
- "quant_summary_path": str(summary_path),
- }
-
-
-def quantization_format(mtq, precision: str) -> dict[str, Any]:
- if precision == "int4":
- return mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG
- if precision == "fp8":
- return mtq.FP8_DEFAULT_CFG
- raise ValueError(f"Unsupported precision: {precision}")
-
-
-def restore_modelopt_quantized_model(
- model,
- modelopt_path: Path,
- *,
- summary_path: Path | None = None,
-) -> tuple[object, dict[str, Any]]:
- import modelopt.torch.opt as mto
- import modelopt.torch.quantization as mtq
-
- model = mto.restore(model, modelopt_path)
- model.eval()
- print_and_store_quant_summary(mtq, model, summary_path, verbose=False)
- return model, {
- "mode": "modelopt-restore",
- "modelopt_model_path": str(modelopt_path),
- "quant_summary_path": str(summary_path) if summary_path is not None else None,
- }
-
-
-def load_quant_model(args: argparse.Namespace) -> QuantModel:
- dtype = torch.float16 if args.dtype == "float16" else torch.float32
- if args.model == "parakeet":
- from rnnt.parakeet_tdt.nemo_backend import DEFAULT_MODEL_ID, MODEL_DIR, load_nemo_model
-
- model_id = args.model_id or DEFAULT_MODEL_ID
- model_dir = args.model_dir or MODEL_DIR
- log(f"loading Parakeet candidate from {model_id}")
- del model_dir
- model = load_nemo_model(
- model_id,
- device=args.device,
- dtype=dtype,
- )
- model.eval()
- return QuantModel(model, int(model.cfg.preprocessor.sample_rate), parakeet_joint_ce_step)
-
- from rnnt.nemotron_asr.nemo_backend import DEFAULT_ATT_CONTEXT_SIZE, DEFAULT_MODEL_ID, load_nemo_model, prompt_id
-
- model_id = args.model_id or DEFAULT_MODEL_ID
- log(f"loading Nemotron candidate from {model_id}")
- model = load_nemo_model(
- model_id,
- device=args.device,
- target_lang=args.target_lang,
- att_context_size=DEFAULT_ATT_CONTEXT_SIZE,
- dtype=dtype,
- )
- model.eval()
- model._din_prompt_id = prompt_id(model, args.target_lang)
- return QuantModel(model, int(model.cfg.preprocessor.sample_rate), nemotron_joint_ce_step)
-
-
-def parakeet_encoder_step(model, item: ValidationItem) -> torch.Tensor:
- samples, lengths = audio_batch([item.audio], model.device)
- features, feature_lengths = model.preprocessor(input_signal=samples, length=lengths)
- compute_dtype = next(model.encoder.parameters()).dtype
- encoded, encoded_lengths = model.encoder(audio_signal=features.to(dtype=compute_dtype), length=feature_lengths)
- states = encoded.transpose(1, 2)
- encoder_mask = torch.arange(states.shape[1], device=states.device)[None] < encoded_lengths[:, None]
- if encoder_mask is not None:
- states = states[:, : int(encoder_mask[0].sum().item())]
- return states
-
-
-def parakeet_joint_ce_step(model, item: ValidationItem, max_steps: int = 256) -> JointLossOutput:
- states = parakeet_encoder_step(model, item)
- valid_length = min(states.shape[1], max_steps)
- if valid_length <= 0:
- raise ValueError(f"Calibration sample {item.id} produced no encoder frames.")
- state = initial_decoder_state(model, states.device)
- token_logits = []
- token_targets = []
- duration_logits = []
- duration_targets = []
- vocab_size = blank_token_id(model) + 1
-
- for time_idx in range(valid_length):
- decoder_input = torch.tensor([[state.last_token]], device=states.device, dtype=torch.long)
- with torch.no_grad():
- decoder_state, (next_hidden, next_cell) = model.decoder.predict(
- y=decoder_input,
- state=(state.hidden, state.cell),
- add_sos=False,
- batch_size=1,
- )
- decoder_state = decoder_state.detach()
- logits = parakeet_joint_logits(model, states[:, time_idx : time_idx + 1], decoder_state).squeeze(0).squeeze(0)
- token_logit = logits[:vocab_size]
- duration_logit = logits[vocab_size:]
- token_target = token_logit.detach().argmax(dim=-1)
- duration_target = duration_logit.detach().argmax(dim=-1)
-
- token_logits.append(token_logit)
- token_targets.append(token_target)
- duration_logits.append(duration_logit)
- duration_targets.append(duration_target)
-
- token = int(token_target.item())
- if token != blank_token_id(model):
- state.last_token = token
- state.hidden = next_hidden.detach()
- state.cell = next_cell.detach()
-
- return JointLossOutput(
- token_logits=torch.stack(token_logits),
- token_targets=torch.stack(token_targets),
- duration_logits=torch.stack(duration_logits),
- duration_targets=torch.stack(duration_targets),
- )
-
-
-def parakeet_joint_logits(model, encoder_states: torch.Tensor, decoder_output: torch.Tensor) -> torch.Tensor:
- joint = model.joint
- if (
- hasattr(joint, "project_encoder")
- and hasattr(joint, "project_prednet")
- and hasattr(joint, "joint_after_projection")
- ):
- enc = joint.project_encoder(encoder_states)
- pred = joint.project_prednet(decoder_output)
- logits = joint.joint_after_projection(enc, pred)
- else:
- logits = joint(encoder_states, decoder_output)
- if logits.dim() == 4 and logits.shape[2] == 1:
- logits = logits.squeeze(2)
- return logits
-
-
-def initial_decoder_state(model, device):
- from types import SimpleNamespace
-
- hidden, cell = model.decoder.initialize_state(torch.zeros(1, 1, joint_encoder_hidden(model), device=device))
- return SimpleNamespace(last_token=blank_token_id(model), hidden=hidden, cell=cell)
-
-
-def blank_token_id(model) -> int:
- return int(model.cfg.decoder.vocab_size)
-
-
-def joint_encoder_hidden(model) -> int:
- prednet = getattr(model.cfg.decoder, "prednet", model.cfg.decoder)
- return int(getattr(model.cfg.joint.jointnet, "encoder_hidden", prednet.pred_hidden))
-
-
-def nemotron_encoder_step(model, item: ValidationItem) -> torch.Tensor:
- samples, lengths = audio_batch([item.audio], model.device)
- features, feature_lengths = model.preprocessor(input_signal=samples, length=lengths)
- compute_dtype = next(model.encoder.parameters()).dtype
- encoded, encoded_lengths = model.encoder(audio_signal=features.to(dtype=compute_dtype), length=feature_lengths)
- states = encoded.transpose(1, 2)
- prompt = torch.zeros(states.shape[0], states.shape[1], 128, device=states.device, dtype=states.dtype)
- prompt[:, :, int(model._din_prompt_id)] = 1.0
- states = model.prompt_kernel(torch.cat([states, prompt], dim=-1))
- return states[:, : int(encoded_lengths[0].item())]
-
-
-def nemotron_joint_ce_step(model, item: ValidationItem, max_steps: int = 256) -> JointLossOutput:
- states = nemotron_encoder_step(model, item)
- valid_length = min(states.shape[1], max_steps)
- if valid_length <= 0:
- raise ValueError(f"Calibration sample {item.id} produced no encoder frames.")
- token_logits = []
- token_targets = []
-
- with torch.no_grad():
- decoder_out, state = model.decoder.predict(y=None, state=None, batch_size=1)
- decoder_out = decoder_out[:, -1:, :].detach()
-
- for time_idx in range(valid_length):
- enc = model.joint.project_encoder(states[:, time_idx : time_idx + 1])
- pred = model.joint.project_prednet(decoder_out)
- logits = model.joint.joint_after_projection(enc, pred).squeeze(0).squeeze(0).squeeze(0)
- token_target = logits.detach().argmax(dim=-1)
- token_logits.append(logits)
- token_targets.append(token_target)
-
- token = int(token_target.item())
- if token != int(model.cfg.decoder.vocab_size):
- token_tensor = torch.tensor([[token]], device=states.device, dtype=torch.long)
- with torch.no_grad():
- decoder_out, state = model.decoder.predict(y=token_tensor, state=state, add_sos=False)
- decoder_out = decoder_out[:, -1:, :].detach()
-
- return JointLossOutput(
- token_logits=torch.stack(token_logits),
- token_targets=torch.stack(token_targets),
- )
-
-
-def audio_batch(audios: list[Audio], device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
- max_len = max(len(audio.samples) for audio in audios)
- batch = torch.zeros(len(audios), max_len, device=device, dtype=torch.float32)
- lengths = torch.empty(len(audios), device=device, dtype=torch.long)
- for index, audio in enumerate(audios):
- samples = torch.from_numpy(audio.samples).to(device=device, dtype=torch.float32)
- batch[index, : samples.numel()] = samples
- lengths[index] = samples.numel()
- return batch, lengths
-
-
-def load_calibration_items(args: argparse.Namespace, sample_rate: int) -> Iterator[ValidationItem]:
- if args.manifest is not None:
- log(f"reading up to {args.calib_samples} samples from manifest {args.manifest}")
- items: Iterable[ValidationItem] = iter_manifest(args.manifest, target_sampling_rate=sample_rate)
- elif args.audio_dir is not None:
- log(f"reading up to {args.calib_samples} samples from audio directory {args.audio_dir}")
- items = iter_audio_dir(args.audio_dir, target_sampling_rate=sample_rate)
- else:
- yield from iter_fleurs_round_robin(args, sample_rate)
- return
-
- for index, item in enumerate(items):
- if index >= args.calib_samples:
- break
- log(f"loaded calibration sample {index + 1}/{args.calib_samples}: {item.id}")
- yield item
-
-
-def iter_fleurs_round_robin(args: argparse.Namespace, sample_rate: int) -> Iterator[ValidationItem]:
- languages = parse_languages(args.languages)
- if not languages:
- return
- per_language = max(1, (args.calib_samples + len(languages) - 1) // len(languages))
- log(
- f"streaming up to {args.calib_samples} samples from FLEURS split={args.split} "
- f"round-robin languages={','.join(languages)}"
- )
- iterators = [
- iter(
- iter_fleurs(
- [language],
- split=args.split,
- max_samples_per_language=per_language,
- target_sampling_rate=sample_rate,
- )
- )
- for language in languages
- ]
-
- yielded = 0
- while iterators and yielded < args.calib_samples:
- remaining = []
- for iterator in iterators:
- if yielded >= args.calib_samples:
- break
- try:
- item = next(iterator)
- except StopIteration:
- continue
- yielded += 1
- log(f"loaded calibration sample {yielded}/{args.calib_samples}: {item.id}")
- yield item
- remaining.append(iterator)
- iterators = remaining
-
-
-def calibration_source(args: argparse.Namespace) -> str:
- if args.manifest is not None:
- return f"manifest {args.manifest}"
- if args.audio_dir is not None:
- return f"audio directory {args.audio_dir}"
- return f"FLEURS split {args.split} languages={args.languages}"
-
-
-def print_and_store_quant_summary(mtq, model, summary_path: Path | None, *, verbose: bool) -> str:
- buffer = io.StringIO()
- with redirect_stdout(buffer):
- mtq.print_quant_summary(model)
- summary = buffer.getvalue()
- if verbose and summary:
- print(summary, end="")
- if summary_path is not None:
- summary_path.write_text(summary, encoding="utf-8")
- return summary
-
-
-def jsonable(value):
- if isinstance(value, dict):
- return {str(key): jsonable(item) for key, item in value.items()}
- if isinstance(value, (list, tuple)):
- return [jsonable(item) for item in value]
- if isinstance(value, torch.Tensor):
- return value.detach().cpu().tolist()
- if isinstance(value, (str, int, float, bool)) or value is None:
- return value
- return repr(value)
-
-
-def parse_args() -> argparse.Namespace:
- parser = argparse.ArgumentParser(description="AutoQuantize the Parakeet or Nemotron encoder with ModelOpt.")
- parser.add_argument("--model", choices=("parakeet", "nemotron"), required=True)
- parser.add_argument("--precision", choices=("int4", "fp8"), required=True)
- parser.add_argument("--output", type=Path, required=True)
- parser.add_argument("--model-id", default=None)
- parser.add_argument("--model-dir", type=Path, default=None)
- parser.add_argument("--device", default="cuda")
- parser.add_argument("--dtype", choices=("float16", "float32"), default="float16")
- parser.add_argument("--effective-bits", type=float, default=None)
- parser.add_argument("--calib-samples", type=int, default=128)
- parser.add_argument("--calib-steps", type=int, default=16)
- parser.add_argument("--score-steps", type=int, default=16)
- parser.add_argument("-v", "--verbose", action="store_true")
- parser.add_argument("--target-lang", default="auto")
- source = parser.add_mutually_exclusive_group()
- source.add_argument("--manifest", type=Path)
- source.add_argument("--audio-dir", type=Path)
- parser.add_argument("--languages", default="en_us,it_it,es_419,de_de")
- parser.add_argument("--split", default=DEFAULT_SPLIT)
- parser.add_argument("--samples-per-language", type=int, default=128)
-
- args = parser.parse_args()
- if args.effective_bits is None:
- args.effective_bits = 8.0 if args.precision == "fp8" else 4.8
- if args.calib_samples <= 0:
- raise ValueError("--calib-samples must be positive.")
- if args.calib_steps <= 0 or args.score_steps <= 0:
- raise ValueError("--calib-steps and --score-steps must be positive.")
- return args
-
-
-def log(message: str) -> None:
- print(f"[quantization] {message}", flush=True)
-
-
-if __name__ == "__main__":
- main()
diff --git a/asr/rnnt/python/tests/__init__.py b/asr/rnnt/python/tests/__init__.py
deleted file mode 100644
index 52a7a9d..0000000
--- a/asr/rnnt/python/tests/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
diff --git a/asr/rnnt/python/tests/support.py b/asr/rnnt/python/tests/support.py
deleted file mode 100644
index 5db6c96..0000000
--- a/asr/rnnt/python/tests/support.py
+++ /dev/null
@@ -1,14 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-import torch
-
-
-def get_device() -> torch.device:
- return torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-
-def inference_dtype(device: torch.device) -> torch.dtype:
- return torch.float16 if device.type == "cuda" else torch.float32
diff --git a/asr/rnnt/python/tests/test_quantization.py b/asr/rnnt/python/tests/test_quantization.py
deleted file mode 100644
index c741c3c..0000000
--- a/asr/rnnt/python/tests/test_quantization.py
+++ /dev/null
@@ -1,30 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-
-def test_quantization_cli_sets_model_defaults(monkeypatch, tmp_path):
- from rnnt.validation import quantization
-
- monkeypatch.setattr(
- "sys.argv",
- ["quantization", "--model", "nemotron", "--precision", "fp8", "--output", str(tmp_path)],
- )
-
- args = quantization.parse_args()
-
- assert args.effective_bits == 8.0
-
-
-def test_quantization_cli_defaults_int4_effective_bits(monkeypatch, tmp_path):
- from rnnt.validation import quantization
-
- monkeypatch.setattr(
- "sys.argv",
- ["quantization", "--model", "parakeet", "--precision", "int4", "--output", str(tmp_path)],
- )
-
- args = quantization.parse_args()
-
- assert args.effective_bits == 4.8
diff --git a/asr/rnnt/python/tests/test_validation_cli.py b/asr/rnnt/python/tests/test_validation_cli.py
deleted file mode 100644
index 8e44b97..0000000
--- a/asr/rnnt/python/tests/test_validation_cli.py
+++ /dev/null
@@ -1,41 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-from pathlib import Path
-
-
-def test_validation_cli_accepts_nemotron_quantized_checkpoint(monkeypatch, tmp_path):
- from rnnt.validation import __main__ as validation
-
- checkpoint = tmp_path / "auto_quantize_model.pt"
- monkeypatch.setattr(
- "sys.argv",
- [
- "validation",
- "--output",
- str(tmp_path / "out"),
- "--asr-model",
- "nemotron",
- "--quantized-modelopt-path",
- str(checkpoint),
- ],
- )
-
- args = validation.parse_args()
-
- assert args.asr_model == "nemotron"
- assert args.quantized_modelopt_path == checkpoint
- assert args.cpp_model_dir == Path("artifacts/nemotron/onnx")
-
-
-def test_validation_cli_defaults_to_parakeet_artifacts(monkeypatch, tmp_path):
- from rnnt.validation import __main__ as validation
-
- monkeypatch.setattr("sys.argv", ["validation", "--output", str(tmp_path / "out")])
-
- args = validation.parse_args()
-
- assert args.asr_model == "parakeet"
- assert args.cpp_model_dir == Path("artifacts/parakeet/onnx")
diff --git a/asr/rnnt/python/tools/__init__.py b/asr/rnnt/python/tools/__init__.py
deleted file mode 100644
index 52a7a9d..0000000
--- a/asr/rnnt/python/tools/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
diff --git a/asr/rnnt/python/tools/plot_validation_wer.py b/asr/rnnt/python/tools/plot_validation_wer.py
deleted file mode 100644
index ca009a0..0000000
--- a/asr/rnnt/python/tools/plot_validation_wer.py
+++ /dev/null
@@ -1,503 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-import argparse
-import csv
-import json
-import math
-import re
-from dataclasses import dataclass
-from pathlib import Path
-
-
-@dataclass(frozen=True)
-class LanguageWer:
- wer: float
- errors: int
- reference_words: int
- count: int
-
-
-@dataclass(frozen=True)
-class ValidationRun:
- name: str
- path: Path
- label: str
- overall: LanguageWer
- by_language: dict[str, LanguageWer]
-
-
-LANGUAGE_LABELS = {
- "de_de": "German",
- "en_us": "English (US)",
- "es_419": "Spanish (LatAm)",
- "fr_fr": "French",
- "it_it": "Italian",
-}
-
-
-def main() -> None:
- args = parse_args()
- runs = discover_runs(args.runs_dir, args.include, args.exclude, args.min_total_count)
- if not runs:
- raise SystemExit(
- f"No validation summaries found in {args.runs_dir} after applying filters. "
- "Try lowering --min-total-count or passing --include."
- )
-
- baseline = select_baseline(runs, args.baseline)
- ordered_runs = [baseline, *sorted((run for run in runs if run != baseline), key=lambda run: run.overall.wer)]
- ordered_runs = with_unique_labels(ordered_runs)
- baseline = ordered_runs[0]
- languages = order_languages(ordered_runs, baseline)
-
- args.output.parent.mkdir(parents=True, exist_ok=True)
- plot_runs(ordered_runs, baseline, languages, args.output)
-
- if args.csv is not None:
- args.csv.parent.mkdir(parents=True, exist_ok=True)
- write_csv(ordered_runs, baseline, languages, args.csv)
-
- print(f"Wrote {args.output}")
- if args.csv is not None:
- print(f"Wrote {args.csv}")
- print(f"Baseline: {baseline.name} ({baseline.overall.wer:.2%} WER, n={baseline.overall.count})")
-
-
-def parse_args() -> argparse.Namespace:
- parser = argparse.ArgumentParser(
- description="Plot per-language WER and relative WER differences from rnnt.validation summaries."
- )
- parser.add_argument(
- "--runs-dir",
- type=Path,
- default=Path("artifacts/validation/runs"),
- help="Directory containing validation run subdirectories with summary.json files.",
- )
- parser.add_argument(
- "--output",
- type=Path,
- default=Path("artifacts/validation/wer_by_language.png"),
- help="Output image path. The extension controls the Matplotlib output format.",
- )
- parser.add_argument(
- "--csv",
- type=Path,
- default=Path("artifacts/validation/wer_by_language.csv"),
- help="Optional CSV output path with the plotted values. Use --csv '' to disable.",
- )
- parser.add_argument(
- "--baseline",
- default=None,
- help="Baseline run name or path. Defaults to the run with the lowest overall WER after filtering.",
- )
- parser.add_argument(
- "--include",
- default=None,
- help="Regex matched against run names. Only matching runs are included.",
- )
- parser.add_argument(
- "--exclude",
- default=None,
- help="Regex matched against run names. Matching runs are excluded.",
- )
- parser.add_argument(
- "--min-total-count",
- type=int,
- default=1000,
- help="Minimum total utterance count required for a run. Defaults to 1000 to hide smoke runs.",
- )
- args = parser.parse_args()
- if args.csv == Path(""):
- args.csv = None
- return args
-
-
-def discover_runs(
- runs_dir: Path,
- include: str | None,
- exclude: str | None,
- min_total_count: int,
-) -> list[ValidationRun]:
- include_re = re.compile(include) if include else None
- exclude_re = re.compile(exclude) if exclude else None
- runs = []
- for summary_path in sorted(runs_dir.glob("*/summary.json")):
- run_name = summary_path.parent.name
- if include_re is not None and include_re.search(run_name) is None:
- continue
- if exclude_re is not None and exclude_re.search(run_name) is not None:
- continue
-
- run = load_run(summary_path)
- if run.overall.count < min_total_count:
- continue
- runs.append(run)
- return runs
-
-
-def load_run(summary_path: Path) -> ValidationRun:
- with summary_path.open("r", encoding="utf-8") as file:
- payload = json.load(file)
-
- summary = payload["summary"]
- overall = LanguageWer(
- wer=float(summary["wer"]),
- errors=int(summary["errors"]),
- reference_words=int(summary["reference_words"]),
- count=int(summary["count"]),
- )
- by_language = {
- language: LanguageWer(
- wer=float(values["wer"]),
- errors=int(values["errors"]),
- reference_words=int(values["reference_words"]),
- count=int(values["count"]),
- )
- for language, values in summary.get("by_language", {}).items()
- }
- return ValidationRun(
- summary_path.parent.name,
- summary_path,
- concise_label(summary_path.parent.name, payload.get("config", {})),
- overall,
- by_language,
- )
-
-
-def select_baseline(runs: list[ValidationRun], baseline: str | None) -> ValidationRun:
- if baseline is None:
- return min(runs, key=lambda run: run.overall.wer)
-
- baseline_path = Path(baseline)
- for run in runs:
- if run.name == baseline or run.path.parent == baseline_path or run.path == baseline_path:
- return run
- choices = ", ".join(run.name for run in runs)
- raise SystemExit(f"Baseline {baseline!r} was not found after filtering. Available runs: {choices}")
-
-
-def order_languages(runs: list[ValidationRun], baseline: ValidationRun) -> list[str]:
- languages = sorted({language for run in runs for language in run.by_language})
- return sorted(
- languages,
- key=lambda language: baseline.by_language.get(language, LanguageWer(0.0, 0, 0, 0)).wer,
- reverse=True,
- )
-
-
-def with_unique_labels(runs: list[ValidationRun]) -> list[ValidationRun]:
- label_counts: dict[str, int] = {}
- unique_runs = []
- for run in runs:
- label_counts[run.label] = label_counts.get(run.label, 0) + 1
- label = run.label if label_counts[run.label] == 1 else f"{run.label} #{label_counts[run.label]}"
- unique_runs.append(ValidationRun(run.name, run.path, label, run.overall, run.by_language))
- return unique_runs
-
-
-def concise_label(name: str, config: dict) -> str:
- backend = config.get("backend")
- dtype = config.get("dtype")
- attention = config.get("attention")
- quantize = config.get("quantize")
-
- if backend == "cpp":
- provider = config.get("cpp_provider") or "cpp"
- return f"C++ {provider.upper()}"
- if quantize and quantize != "none":
- bits = config.get("auto_quant_effective_bits")
- bits_suffix = f" {bits:g}b" if isinstance(bits, (int, float)) else ""
- return f"{quantize.replace('-', ' ').title().replace('Int4', 'INT4')}{bits_suffix}"
- if attention == "rel_pos_local_attn":
- context = str(config.get("att_context_size") or "").replace(",", "/")
- context_suffix = f" {context}" if context else ""
- dtype_suffix = f" {dtype}" if dtype else ""
- return f"Local attn{context_suffix}{dtype_suffix}"
- if dtype:
- return f"Full attn {dtype}"
- return name.replace("_", " ")
-
-
-def plot_runs(
- runs: list[ValidationRun],
- baseline: ValidationRun,
- languages: list[str],
- output: Path,
-) -> None:
- try:
- import matplotlib.pyplot as plt
- except ImportError as error:
- raise SystemExit(
- "Plotting requires matplotlib. Install it with `python -m pip install -e .[validation]`."
- ) from error
-
- plt.rcParams.update(
- {
- "axes.spines.top": False,
- "axes.spines.right": False,
- "axes.titleweight": "bold",
- "font.size": 10,
- "legend.frameon": False,
- }
- )
-
- width = max(11.5, 1.1 * len(languages) + 4.0)
- height = 7.0 if len(runs) > 1 else 4.8
- figure, axes = plt.subplots(
- 2 if len(runs) > 1 else 1,
- 1,
- figsize=(width, height),
- gridspec_kw={"height_ratios": [1.2, 1.0]} if len(runs) > 1 else None,
- )
- figure.subplots_adjust(top=0.82, hspace=0.34, left=0.13, right=0.98, bottom=0.08)
- if len(runs) == 1:
- wer_axis = axes
- delta_axis = None
- else:
- wer_axis, delta_axis = axes
-
- y_positions = list(range(len(languages)))
- bar_height = min(0.72 / len(runs), 0.18)
- offsets = centered_offsets(len(runs), bar_height)
- colors = palette(plt, runs, baseline)
-
- for index, run in enumerate(runs):
- values = [percent(run.by_language.get(language)) for language in languages]
- wer_axis.barh(
- [position + offsets[index] for position in y_positions],
- values,
- height=bar_height,
- label=legend_label(run, baseline),
- color=colors[run.name],
- alpha=0.92,
- )
-
- wer_axis.set_title("WER by language")
- wer_axis.set_xlabel("WER (%)")
- wer_axis.set_yticks(y_positions, [language_label(language) for language in languages])
- wer_axis.invert_yaxis()
- wer_axis.grid(axis="x", linestyle="-", alpha=0.18)
- wer_axis.set_axisbelow(True)
- annotate_absolute_values(wer_axis, runs, languages, y_positions, offsets, colors)
-
- if delta_axis is not None:
- comparison_runs = [run for run in runs if run != baseline]
- delta_height = min(0.72 / max(len(comparison_runs), 1), 0.22)
- delta_offsets = centered_offsets(len(comparison_runs), delta_height)
- for index, run in enumerate(comparison_runs):
- deltas = [
- relative_delta_percent(run.by_language.get(language), baseline.by_language.get(language))
- for language in languages
- ]
- delta_axis.barh(
- [position + delta_offsets[index] for position in y_positions],
- deltas,
- height=delta_height,
- label=legend_label(run, baseline),
- color=colors[run.name],
- alpha=0.92,
- )
-
- delta_axis.axvline(0.0, color="#303030", linewidth=1.0)
- delta_axis.set_title(f"Relative WER change vs {baseline.label}")
- delta_axis.set_xlabel("Change in WER (%)")
- delta_axis.set_yticks(y_positions, [language_label(language) for language in languages])
- delta_axis.invert_yaxis()
- delta_axis.grid(axis="x", linestyle="-", alpha=0.18)
- delta_axis.set_axisbelow(True)
- annotate_delta_values(delta_axis, comparison_runs, baseline, languages, y_positions, delta_offsets)
-
- handles, labels = wer_axis.get_legend_handles_labels()
- figure.legend(
- handles,
- labels,
- loc="upper center",
- bbox_to_anchor=(0.5, 0.9),
- ncols=min(len(runs), 4),
- )
- figure.suptitle("Validation WER Comparison", fontsize=16, fontweight="bold", y=0.98)
- figure.text(
- 0.5,
- 0.94,
- f"Baseline: {baseline.label} ({baseline.overall.wer:.2%} overall WER, n={baseline.overall.count}). Lower is better.",
- ha="center",
- fontsize=10,
- color="#555555",
- )
- figure.savefig(output, dpi=180, bbox_inches="tight")
- plt.close(figure)
-
-
-def centered_offsets(count: int, width: float) -> list[float]:
- center = (count - 1) / 2
- return [(index - center) * width for index in range(count)]
-
-
-def language_label(language: str) -> str:
- return LANGUAGE_LABELS.get(language, language)
-
-
-def palette(plt, runs: list[ValidationRun], baseline: ValidationRun) -> dict[str, str]:
- color_cycle = list(plt.get_cmap("Set2").colors) + list(plt.get_cmap("tab10").colors)
- colors = {baseline.name: "#2f3437"}
- color_index = 0
- for run in runs:
- if run == baseline:
- continue
- colors[run.name] = color_cycle[color_index % len(color_cycle)]
- color_index += 1
- return colors
-
-
-def legend_label(run: ValidationRun, baseline: ValidationRun) -> str:
- prefix = "Baseline: " if run == baseline else ""
- return f"{prefix}{run.label} ({run.overall.wer:.2%})"
-
-
-def annotate_absolute_values(
- axis,
- runs: list[ValidationRun],
- languages: list[str],
- y_positions: list[int],
- offsets: list[float],
- colors: dict[str, str],
-) -> None:
- max_value = max(
- percent(run.by_language.get(language))
- for run in runs
- for language in languages
- if run.by_language.get(language) is not None
- )
- axis.set_xlim(0, max_value * 1.16)
- for run_index, run in enumerate(runs):
- for y_position, language in zip(y_positions, languages, strict=True):
- value = percent(run.by_language.get(language))
- if math.isnan(value):
- continue
- axis.text(
- value + max_value * 0.012,
- y_position + offsets[run_index],
- f"{value:.1f}",
- va="center",
- ha="left",
- fontsize=8,
- color=colors[run.name],
- )
-
-
-def annotate_delta_values(
- axis,
- runs: list[ValidationRun],
- baseline: ValidationRun,
- languages: list[str],
- y_positions: list[int],
- offsets: list[float],
-) -> None:
- visible_values = []
- for run in runs:
- for language in languages:
- value = relative_delta_percent(run.by_language.get(language), baseline.by_language.get(language))
- if not math.isnan(value):
- visible_values.append(value)
-
- if not visible_values:
- return
-
- min_value = min(visible_values)
- max_value = max(visible_values)
- max_abs = max(abs(min_value), abs(max_value), 1.0)
- axis.set_xlim(min(min_value * 1.2, -0.12 * max_abs), max(max_value * 1.2, 0.12 * max_abs))
- axis.axvspan(axis.get_xlim()[0], 0, color="#2ca25f", alpha=0.06, zorder=0)
- axis.axvspan(0, axis.get_xlim()[1], color="#de2d26", alpha=0.05, zorder=0)
-
- label_offset = max_abs * 0.025
- for run_index, run in enumerate(runs):
- for y_position, language in zip(y_positions, languages, strict=True):
- value = relative_delta_percent(run.by_language.get(language), baseline.by_language.get(language))
- if math.isnan(value):
- continue
- axis.text(
- value + (label_offset if value >= 0 else -label_offset),
- y_position + offsets[run_index],
- f"{value:+.1f}%",
- va="center",
- ha="left" if value >= 0 else "right",
- fontsize=8,
- color="#404040",
- )
-
-
-def percent(value: LanguageWer | None) -> float:
- if value is None:
- return math.nan
- return 100.0 * value.wer
-
-
-def relative_delta_percent(value: LanguageWer | None, baseline: LanguageWer | None) -> float:
- if value is None or baseline is None or baseline.wer == 0.0:
- return math.nan
- return 100.0 * (value.wer - baseline.wer) / baseline.wer
-
-
-def write_csv(
- runs: list[ValidationRun],
- baseline: ValidationRun,
- languages: list[str],
- csv_path: Path,
-) -> None:
- fields = [
- "run",
- "language",
- "wer",
- "wer_percent",
- "errors",
- "reference_words",
- "count",
- "baseline_run",
- "baseline_wer",
- "absolute_delta",
- "relative_delta_percent",
- ]
- with csv_path.open("w", newline="", encoding="utf-8") as file:
- writer = csv.DictWriter(file, fieldnames=fields)
- writer.writeheader()
- for run in runs:
- write_csv_row(writer, run, "__overall__", run.overall, baseline.name, baseline.overall)
- for language in languages:
- value = run.by_language.get(language)
- baseline_value = baseline.by_language.get(language)
- if value is None:
- continue
- write_csv_row(writer, run, language, value, baseline.name, baseline_value)
-
-
-def write_csv_row(
- writer: csv.DictWriter,
- run: ValidationRun,
- language: str,
- value: LanguageWer,
- baseline_name: str,
- baseline: LanguageWer | None,
-) -> None:
- writer.writerow(
- {
- "run": run.name,
- "language": language,
- "wer": value.wer,
- "wer_percent": percent(value),
- "errors": value.errors,
- "reference_words": value.reference_words,
- "count": value.count,
- "baseline_run": baseline_name,
- "baseline_wer": baseline.wer if baseline is not None else "",
- "absolute_delta": value.wer - baseline.wer if baseline is not None else "",
- "relative_delta_percent": relative_delta_percent(value, baseline),
- }
- )
-
-
-if __name__ == "__main__":
- main()
diff --git a/asr/whisper/CMakeLists.txt b/asr/whisper/CMakeLists.txt
index 2e6e6fc..0f4caf6 100644
--- a/asr/whisper/CMakeLists.txt
+++ b/asr/whisper/CMakeLists.txt
@@ -1,5 +1,4 @@
-# Whisper speech-to-text. Follows the asr/rnnt layering: a processing library
-# (din_asr_whisper) plus a thin CLI (din_asr_whisper)
+# Whisper speech-to-text: a processing library plus a thin CLI.
# Optional CUDA greedy-argmax kernel for the TRT-RTX decode path. When CUDA is
# unavailable the pipeline falls back to a pinned-host CPU argmax (DIN_WHISPER_CUDA
@@ -23,6 +22,7 @@ target_include_directories(din_asr_whisper
target_link_libraries(din_asr_whisper
PUBLIC
+ din_asr_common
din_common_ort
din_common_io
nlohmann_json::nlohmann_json
diff --git a/asr/whisper/whisper.h b/asr/whisper/whisper.h
index f4e721b..e5d06f5 100644
--- a/asr/whisper/whisper.h
+++ b/asr/whisper/whisper.h
@@ -12,17 +12,13 @@
#include
#include
+#include "asr_common.h"
#include "audio.h"
#include "ort_session.h"
#include "tokenizer.h"
#include
-// --- timing helper -------------------------------
-
-inline double SecondsSince(std::chrono::steady_clock::time_point start)
-{
- return std::chrono::duration(std::chrono::steady_clock::now() - start).count();
-}
+using din::asr::common::SecondsSince;
namespace din::asr::whisper
{
diff --git a/asr/rnnt/python/tools/nemo_preprocessor_export.py b/common/model_export/nemo_export.py
similarity index 76%
rename from asr/rnnt/python/tools/nemo_preprocessor_export.py
rename to common/model_export/nemo_export.py
index ff8bff0..864d903 100644
--- a/asr/rnnt/python/tools/nemo_preprocessor_export.py
+++ b/common/model_export/nemo_export.py
@@ -106,3 +106,31 @@ def normalize_batch_without_item_guard(x, seq_len, normalize_type):
normalized = (x - x_mean.unsqueeze(2)) / x_std.unsqueeze(2)
normalized = normalized.masked_fill(~valid_mask.unsqueeze(1), 0.0)
return normalized, x_mean, x_std
+
+
+def load_nemo_model(
+ model_id: str, *, device: str | torch.device = "cpu", dtype: torch.dtype = torch.float32
+) -> nn.Module:
+ target_device = torch.device(device)
+ return _load_model(model_id, target_device).to(device=target_device, dtype=dtype).eval()
+
+
+def _load_model(model_id: str, device: torch.device) -> nn.Module:
+ try:
+ from nemo.collections import asr as nemo_asr
+ except ImportError as exc:
+ raise RuntimeError("NeMo export requires the optional `nemo_toolkit[asr]` dependency.") from exc
+
+ model_path = Path(model_id)
+ if model_path.is_file() and model_path.suffix == ".nemo":
+ return nemo_asr.models.ASRModel.restore_from(str(model_path), map_location=device)
+ return nemo_asr.models.ASRModel.from_pretrained(model_name=model_id, map_location=device)
+
+
+def configure_encoder_sdpa(encoder: nn.Module) -> None:
+ for module in [encoder, *getattr(encoder, "layers", [])]:
+ if hasattr(module, "use_pytorch_sdpa"):
+ module.use_pytorch_sdpa = True
+ self_attn = getattr(module, "self_attn", None)
+ if self_attn is not None and hasattr(self_attn, "use_pytorch_sdpa"):
+ self_attn.use_pytorch_sdpa = True
diff --git a/common/model_export/requirements-nemo.txt b/common/model_export/requirements-nemo.txt
new file mode 100644
index 0000000..da0f056
--- /dev/null
+++ b/common/model_export/requirements-nemo.txt
@@ -0,0 +1,5 @@
+torch
+onnx
+onnxscript
+nemo_toolkit[asr] @ git+https://github.com/NVIDIA-NeMo/Speech.git@v3.0.0
+jiwer
diff --git a/common/model_export/validate_asr.py b/common/model_export/validate_asr.py
new file mode 100644
index 0000000..6a10761
--- /dev/null
+++ b/common/model_export/validate_asr.py
@@ -0,0 +1,62 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""One-file ASR validation through the native CLI; no Python runtime bindings."""
+
+import argparse
+import contextlib
+import json
+import math
+import subprocess
+import sys
+from pathlib import Path
+
+
+def validation_parser(model_id: str) -> argparse.ArgumentParser:
+ parser = argparse.ArgumentParser(description="Compare the C++ ASR CLI with a NeMo reference on real audio.")
+ parser.add_argument("--audio", type=Path, required=True, help="Real audio file (16 kHz mono WAV recommended).")
+ parser.add_argument("--cli", type=Path, required=True)
+ parser.add_argument("--model-dir", type=Path, required=True)
+ parser.add_argument("--model", default=model_id, help="NeMo model ID or local .nemo checkpoint used for export.")
+ parser.add_argument("--device", default="cpu", help="Device for the NeMo reference.")
+ parser.add_argument("--provider", choices=("cpu", "trt-rtx"), default="cpu")
+ parser.add_argument("--truth", type=Path, help="Optional human reference transcript.")
+ parser.add_argument(
+ "--max-wer", type=float, help="Fail above this WER; uses --truth when supplied, otherwise NeMo."
+ )
+ return parser
+
+
+def validate(args: argparse.Namespace, reference, cli_args=()) -> None:
+ # Reject missing inputs before loading a large reference model.
+ for path in (args.audio, args.cli):
+ if not path.is_file():
+ raise FileNotFoundError(path)
+ if args.max_wer is not None and (not math.isfinite(args.max_wer) or args.max_wer < 0):
+ raise ValueError("--max-wer must be finite and nonnegative.")
+ truth = args.truth.read_text(encoding="utf-8").strip() if args.truth else None
+ command = [
+ str(args.cli.resolve()),
+ str(args.audio.resolve()),
+ "--model-dir",
+ str(args.model_dir.resolve()),
+ "--provider",
+ args.provider,
+ "--json",
+ *cli_args,
+ ]
+ result = subprocess.run(command, check=True, stdout=subprocess.PIPE, text=True, encoding="utf-8")
+ hypothesis = json.loads(result.stdout)["transcription"]
+ with contextlib.redirect_stdout(sys.stderr):
+ expected = reference(args)
+ from jiwer import Compose, RemoveMultipleSpaces, RemovePunctuation, Strip, ToLowerCase, wer
+
+ normalize = Compose([ToLowerCase(), RemovePunctuation(), RemoveMultipleSpaces(), Strip()])
+ score = wer(normalize(expected), normalize(hypothesis))
+ output = {"reference": expected, "transcription": hypothesis, "wer_against_nemo": score}
+ if truth is not None:
+ output["truth"] = truth
+ score = output["wer_against_truth"] = wer(normalize(truth), normalize(hypothesis))
+ print(json.dumps(output, ensure_ascii=False, indent=2))
+ if args.max_wer is not None and score > args.max_wer:
+ raise SystemExit(f"WER {score:.4f} exceeds {args.max_wer:.4f}")
diff --git a/pyproject.toml b/pyproject.toml
index 54f617f..0a0d4c8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
[build-system]
-requires = ["scikit-build-core>=0.10", "nanobind", "setuptools>=69", "wheel"]
-build-backend = "scikit_build_core.build"
+requires = ["setuptools>=69", "wheel"]
+build-backend = "setuptools.build_meta"
[project]
name = "din-deploy"
@@ -20,8 +20,6 @@ dependencies = [
"onnxscript",
"safetensors",
"tqdm",
- "nanobind",
- "matplotlib",
"Pillow",
]
@@ -33,21 +31,10 @@ dev = [
parakeet = [
"nemo_toolkit[asr] @ git+https://github.com/NVIDIA-NeMo/Speech.git@v3.0.0",
"jiwer",
- "librosa>=0.10.2",
- "numba>=0.59",
- "nvidia-modelopt",
- "soundfile",
]
nemotron = [
"nemo_toolkit[asr] @ git+https://github.com/NVIDIA-NeMo/Speech.git@v3.0.0",
- "librosa>=0.10.2",
- "numba>=0.59",
- "PyYAML",
- "sentencepiece",
- "soundfile",
-]
-nemotron-reference = [
- "nemo_toolkit[asr] @ git+https://github.com/NVIDIA-NeMo/Speech.git@v3.0.0",
+ "jiwer",
]
sam2 = [
"av",
@@ -65,29 +52,14 @@ flux = [
"transformers",
]
-[project.scripts]
-parakeet-tdt = "rnnt.parakeet_tdt.cli:main"
-nemotron-asr = "rnnt.nemotron_asr.cli:main"
-rnnt-quantize-encoder = "rnnt.validation.quantization:main"
-
-[tool.scikit-build]
-wheel.packages = [
- "asr/rnnt/python/rnnt",
-]
-wheel.py-api = "cp312"
-
-[tool.scikit-build.cmake.define]
-DIN_ENABLE_PARAKEET_PYTHON_BINDINGS = "ON"
-DIN_ENABLE_NEMOTRON_PYTHON_BINDINGS = "ON"
-
-[tool.pytest.ini_options]
-testpaths = ["asr/rnnt/python/tests"]
-pythonpath = ["asr/rnnt/python", "asr/rnnt/python/tests", "vision/sam2/python"]
+# Python is used for model preparation; inference is provided by the C++ CLIs.
+[tool.setuptools]
+packages = []
[tool.ruff]
line-length = 120
target-version = "py312"
-src = ["asr/rnnt/python/rnnt", "asr/rnnt/python/tools", "asr/rnnt/python/tests", "vision/sam2/python"]
+src = ["asr/parakeet/model_export", "asr/nemotron/model_export", "common/model_export", "vision/sam2/python"]
[tool.ruff.lint]
select = [
@@ -109,4 +81,5 @@ indent-style = "space"
line-ending = "auto"
[tool.ruff.lint.per-file-ignores]
-"asr/rnnt/python/tools/*.py" = ["E402"]
+"asr/parakeet/model_export/*.py" = ["E402"]
+"asr/nemotron/model_export/*.py" = ["E402"]
diff --git a/tests/asr_cli/CMakeLists.txt b/tests/asr_cli/CMakeLists.txt
new file mode 100644
index 0000000..dcdf6a0
--- /dev/null
+++ b/tests/asr_cli/CMakeLists.txt
@@ -0,0 +1,6 @@
+add_executable(din_asr_cli_test test_asr_cli.cpp)
+target_link_libraries(din_asr_cli_test PRIVATE din_asr_common argparse nlohmann_json::nlohmann_json)
+if(DIN_BUILD_TESTING)
+ add_test(NAME din_asr_cli_test COMMAND din_asr_cli_test)
+ set_tests_properties(din_asr_cli_test PROPERTIES LABELS "unit")
+endif()
diff --git a/tests/asr_cli/test_asr_cli.cpp b/tests/asr_cli/test_asr_cli.cpp
new file mode 100644
index 0000000..f52ee7b
--- /dev/null
+++ b/tests/asr_cli/test_asr_cli.cpp
@@ -0,0 +1,91 @@
+// SPDX-License-Identifier: Apache-2.0
+#include "cli.h"
+
+namespace
+{
+struct Config
+{
+ std::filesystem::path model_dir = "models";
+ std::string provider = "cpu";
+ std::filesystem::path ep_cache_dir = "cache";
+ std::filesystem::path ep_context_dir = "context";
+};
+struct Result
+{
+ std::string text = "Hello 中";
+ double audio_seconds = 1.0;
+ double transcribe_seconds = 0.25;
+ double encode_seconds = 0.2;
+ double greedy_seconds = 0.05;
+};
+struct Pipeline
+{
+ explicit Pipeline(Config) {}
+ Result TranscribeFile(const std::filesystem::path& path)
+ {
+ if (path == "fail.wav")
+ throw std::runtime_error("inference failed");
+ return {};
+ }
+ void Print(std::ostream& out, const Result& result, const din::asr::common::TranscriptionOptions&) const
+ {
+ out << result.text << '\n';
+ }
+};
+struct Output
+{
+ int code;
+ std::string out;
+ std::string error;
+};
+Output run(std::vector args)
+{
+ std::vector argv;
+ for (auto& arg : args)
+ argv.push_back(arg.data());
+ std::ostringstream out, error;
+ auto* previous_out = std::cout.rdbuf(out.rdbuf());
+ auto* previous_error = std::cerr.rdbuf(error.rdbuf());
+ const din::asr::common::CliSpec spec{.fallback_name = "test_asr",
+ .artifact_label = "Test",
+ .default_provider = "cpu",
+ .timestamps_modes = "none|token",
+ .allowed_timestamps = {"none", "token"}};
+ const int code = din::asr::common::RunCli(static_cast(argv.size()), argv.data(), spec);
+ std::cout.rdbuf(previous_out);
+ std::cerr.rdbuf(previous_error);
+ return {code, out.str(), error.str()};
+}
+void require(bool condition, const char* message)
+{
+ if (!condition)
+ throw std::runtime_error(message);
+}
+} // namespace
+
+int main()
+{
+ try
+ {
+ const auto json = run({"test_asr", "audio.wav", "--json"});
+ require(json.code == 0, "JSON command failed");
+ const auto document = nlohmann::json::parse(json.out);
+ require(document.at("transcription") == "Hello 中", "JSON lost Unicode text");
+ require(document.at("audio_seconds") == 1.0, "JSON lost audio duration");
+ require(json.error.find("provider: cpu") != std::string::npos, "JSON timings must go to stderr");
+ const auto plain = run({"test_asr", "audio.wav"});
+ require(plain.code == 0 && plain.out.starts_with("Hello 中\n"), "Plain transcription changed");
+ require(plain.out.find("provider: cpu") != std::string::npos, "Plain timings changed");
+ const auto invalid = run({"test_asr", "audio.wav", "--timestamps", "invalid"});
+ require(invalid.code != 0 && invalid.out.empty(), "Invalid options must fail without a transcript");
+ const auto failure = run({"test_asr", "fail.wav", "--json"});
+ require(failure.code != 0 && failure.out.empty(), "Inference failure must not print a JSON transcript");
+ std::cout << "ASR CLI checks passed\n";
+ return 0;
+ }
+ catch (const std::exception& error)
+ {
+ std::cerr << error.what() << '\n';
+ return 1;
+ }
+}
diff --git a/tests/model_export/test_validate_asr.py b/tests/model_export/test_validate_asr.py
new file mode 100644
index 0000000..f64a9b2
--- /dev/null
+++ b/tests/model_export/test_validate_asr.py
@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+import importlib.util
+import json
+import subprocess
+from pathlib import Path
+
+import pytest
+
+SPEC = importlib.util.spec_from_file_location(
+ "validate_asr", Path(__file__).resolve().parents[2] / "common/model_export/validate_asr.py"
+)
+validation = importlib.util.module_from_spec(SPEC)
+SPEC.loader.exec_module(validation)
+
+
+@pytest.fixture
+def args(tmp_path):
+ audio = tmp_path / "recording.wav"
+ audio.touch()
+ cli = tmp_path / "asr cli"
+ cli.touch()
+ return validation.validation_parser("test-model").parse_args(
+ ["--audio", str(audio), "--cli", str(cli), "--model-dir", str(tmp_path)]
+ )
+
+
+def cli_result(monkeypatch, text):
+ def run(command, **kwargs):
+ assert command[0].endswith("asr cli")
+ assert "--json" in command
+ assert kwargs["check"] is True
+ return subprocess.CompletedProcess(command, 0, json.dumps({"transcription": text}))
+
+ monkeypatch.setattr(validation.subprocess, "run", run)
+
+
+def test_reference_comparison_and_logging(args, monkeypatch, capsys):
+ cli_result(monkeypatch, "Hello, WORLD!")
+
+ def reference(_):
+ print("reference model diagnostics")
+ return "hello world"
+
+ validation.validate(args, reference)
+ out = capsys.readouterr()
+ assert json.loads(out.out)["wer_against_nemo"] == 0
+ assert "reference model diagnostics" in out.err
+
+
+def test_truth_controls_threshold(args, monkeypatch, capsys):
+ cli_result(monkeypatch, "hello world")
+ args.truth = args.audio.with_suffix(".txt")
+ args.truth.write_text("hello world", encoding="utf-8")
+ args.max_wer = 0
+ validation.validate(args, lambda _: "something different")
+ report = json.loads(capsys.readouterr().out)
+ assert report["wer_against_truth"] == 0
+ assert report["wer_against_nemo"] == 1
+
+
+def test_wer_threshold_failure(args, monkeypatch):
+ cli_result(monkeypatch, "wrong")
+ args.max_wer = 0.1
+ with pytest.raises(SystemExit, match="exceeds"):
+ validation.validate(args, lambda _: "correct")
+
+
+def test_inference_failure_is_not_scored(args, monkeypatch):
+ def run(*_, **__):
+ raise subprocess.CalledProcessError(1, ["asr"])
+
+ monkeypatch.setattr(validation.subprocess, "run", run)
+ with pytest.raises(subprocess.CalledProcessError):
+ validation.validate(args, lambda _: pytest.fail("Must not load reference after CLI failure"))
+
+
+@pytest.mark.parametrize("limit", [-1, float("nan"), float("inf")])
+def test_invalid_threshold_rejected_before_inference(args, monkeypatch, limit):
+ args.max_wer = limit
+ monkeypatch.setattr(validation.subprocess, "run", lambda *_a, **_k: pytest.fail("Must not run inference"))
+ with pytest.raises(ValueError, match="finite and nonnegative"):
+ validation.validate(args, lambda _: "unused")