diff --git a/.github/workflows/export-nemotron-3.5-asr-streaming-0.6b.yaml b/.github/workflows/export-nemotron-3.5-asr-streaming-0.6b.yaml new file mode 100644 index 0000000000..8f9585e3f7 --- /dev/null +++ b/.github/workflows/export-nemotron-3.5-asr-streaming-0.6b.yaml @@ -0,0 +1,185 @@ +name: export-nemotron-3-5-asr-streaming-06b + +on: + push: + branches: + - export-nemotron-3-5-asr-streaming + workflow_dispatch: + +concurrency: + group: export-nemotron-3-5-asr-streaming-to-onnx-${{ github.ref }} + cancel-in-progress: true + +jobs: + export-nemotron-3-5-asr-streaming-0-6b-to-onnx: + if: github.repository_owner == 'k2-fsa' || github.repository_owner == 'csukuangfj' + name: nemotron-3-5-asr-streaming-0-6b-to-onnx + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: [macos-latest] + python-version: ["3.10"] + + steps: + - uses: actions/checkout@v4 + + - name: Setup Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install NeMo + shell: bash + run: | + # nemotron-3.5-asr-streaming-0.6b requires EncDecRNNTBPEModelWithPrompt, + # which is not in the stable nemo-toolkit release yet; the model card + # instructs installing NeMo from main. + BRANCH='main' + pip install Cython packaging + pip install "nemo_toolkit[asr] @ git+https://github.com/NVIDIA/NeMo.git@$BRANCH" + pip install onnxruntime ipython sentencepiece + pip install kaldi-native-fbank + pip install soundfile librosa + + - name: Run + shell: bash + run: | + cd scripts/nemo/nemotron-3.5-asr-streaming-0.6b + + python3 ./export_onnx.py + + ls -lh + echo "---" + + ls -lh */ + + echo "---" + + - name: Collect results + shell: bash + run: | + src=scripts/nemo/nemotron-3.5-asr-streaming-0.6b + + for chunk in 80 160 560 1120; do + d=sherpa-onnx-nemotron-3.5-asr-streaming-0.6b-${chunk}ms-2026-06-11 + mkdir -p $d + + cp -av $src/$chunk/encoder.onnx $d/ + cp -av $src/$chunk/encoder.data $d/ + cp -av $src/$chunk/decoder.onnx $d/ + cp -av $src/$chunk/joiner.onnx $d/ + cp -av $src/tokens.txt $d/ + cat >$d/README.md <$d/README.md < { sherpa_onnx.initBindings(); _recognizer = await createOnlineRecognizer(); _stream = _recognizer?.createStream(); + // Multilingual Nemotron: _stream?.setOption(key: 'language', value: 'ja'); _isInitialized = true; } diff --git a/flutter/sherpa_onnx/lib/src/online_stream.dart b/flutter/sherpa_onnx/lib/src/online_stream.dart index d07384b945..0c5904761e 100644 --- a/flutter/sherpa_onnx/lib/src/online_stream.dart +++ b/flutter/sherpa_onnx/lib/src/online_stream.dart @@ -67,5 +67,25 @@ class OnlineStream { SherpaOnnxBindings.onlineStreamInputFinished?.call(ptr); } + /// Set a string option on the underlying stream. + /// + /// For multilingual Nemotron models, use `key: 'language'` with values such + /// as `ja` or `auto`. + void setOption({required String key, required String value}) { + if (SherpaOnnxBindings.onlineStreamSetOption == null) { + throw Exception("Please initialize sherpa-onnx first"); + } + + if (ptr == nullptr) { + return; + } + + final pKey = key.toNativeUtf8(); + final pValue = value.toNativeUtf8(); + SherpaOnnxBindings.onlineStreamSetOption?.call(ptr, pKey, pValue); + calloc.free(pKey); + calloc.free(pValue); + } + Pointer ptr; } diff --git a/flutter/sherpa_onnx/lib/src/sherpa_onnx_bindings.dart b/flutter/sherpa_onnx/lib/src/sherpa_onnx_bindings.dart index e093c9dde5..6ee7b7fac3 100644 --- a/flutter/sherpa_onnx/lib/src/sherpa_onnx_bindings.dart +++ b/flutter/sherpa_onnx/lib/src/sherpa_onnx_bindings.dart @@ -1795,6 +1795,20 @@ typedef OnlineStreamInputFinishedNative = typedef OnlineStreamInputFinished = void Function(Pointer); +typedef OnlineStreamSetOptionNative = + Void Function( + Pointer, + Pointer, + Pointer, + ); + +typedef OnlineStreamSetOption = + void Function( + Pointer, + Pointer, + Pointer, + ); + typedef SherpaOnnxSpeakerEmbeddingExtractorIsReadyNative = Int32 Function( Pointer, @@ -2043,6 +2057,8 @@ class SherpaOnnxBindings { static OnlineStreamInputFinished? onlineStreamInputFinished; + static OnlineStreamSetOption? onlineStreamSetOption; + static SherpaOnnxSpeakerEmbeddingExtractorIsReady? speakerEmbeddingExtractorIsReady; @@ -2749,6 +2765,12 @@ class SherpaOnnxBindings { ) .asFunction(); + onlineStreamSetOption ??= dynamicLibrary + .lookup>( + 'SherpaOnnxOnlineStreamSetOption', + ) + .asFunction(); + speakerEmbeddingExtractorIsReady ??= dynamicLibrary .lookup< NativeFunction diff --git a/scripts/nemo/nemotron-3.5-asr-streaming-0.6b/README.md b/scripts/nemo/nemotron-3.5-asr-streaming-0.6b/README.md new file mode 100644 index 0000000000..39314f4f6e --- /dev/null +++ b/scripts/nemo/nemotron-3.5-asr-streaming-0.6b/README.md @@ -0,0 +1,64 @@ +# Multilingual Nemotron-3.5 Streaming ASR + +This directory exports the NVIDIA NeMo model +`nvidia/nemotron-3.5-asr-streaming-0.6b` to sherpa-onnx streaming transducer +packages. + +The exporter writes the same package layout as +`nvidia/nemotron-speech-streaming-en-0.6b`: + +- `encoder.onnx` +- `encoder.data` +- `decoder.onnx` +- `joiner.onnx` +- int8 variants for all three ONNX graphs +- `tokens.txt` converted from the model's SentencePiece tokenizer + +The encoder metadata contains `prompt_dictionary` and `auto_prompt_id`. Users +set the per-stream language as a string; the numerical prompt id is internal. +An empty language string and `auto` use the model's auto-detect prompt. + +## Export + +As of June 2026, stable NeMo releases lack `EncDecRNNTBPEModelWithPrompt`, +so install NeMo from git main. + +```bash +pip install Cython packaging +pip install "nemo_toolkit[asr] @ git+https://github.com/NVIDIA/NeMo.git@main" +pip install onnxruntime ipython sentencepiece +pip install kaldi-native-fbank +pip install soundfile librosa + +python3 ./export_onnx.py +``` + +The script exports 80ms, 160ms, 560ms, and 1120ms chunk sizes. + +## Decode + +Forced language: + +```bash +./build/bin/sherpa-onnx \ + --encoder=./sherpa-onnx-nemotron-3.5-asr-streaming-0.6b-560ms-int8-2026-06-11/encoder.int8.onnx \ + --decoder=./sherpa-onnx-nemotron-3.5-asr-streaming-0.6b-560ms-int8-2026-06-11/decoder.int8.onnx \ + --joiner=./sherpa-onnx-nemotron-3.5-asr-streaming-0.6b-560ms-int8-2026-06-11/joiner.int8.onnx \ + --tokens=./sherpa-onnx-nemotron-3.5-asr-streaming-0.6b-560ms-int8-2026-06-11/tokens.txt \ + --language=ja \ + ./sherpa-onnx-nemotron-3.5-asr-streaming-0.6b-560ms-int8-2026-06-11/test_wavs/ja.wav +``` + +Auto language: + +```bash +./build/bin/sherpa-onnx \ + --encoder=./sherpa-onnx-nemotron-3.5-asr-streaming-0.6b-560ms-int8-2026-06-11/encoder.int8.onnx \ + --decoder=./sherpa-onnx-nemotron-3.5-asr-streaming-0.6b-560ms-int8-2026-06-11/decoder.int8.onnx \ + --joiner=./sherpa-onnx-nemotron-3.5-asr-streaming-0.6b-560ms-int8-2026-06-11/joiner.int8.onnx \ + --tokens=./sherpa-onnx-nemotron-3.5-asr-streaming-0.6b-560ms-int8-2026-06-11/tokens.txt \ + --language=auto \ + ./sherpa-onnx-nemotron-3.5-asr-streaming-0.6b-560ms-int8-2026-06-11/test_wavs/ja.wav +``` + +The same auto behavior is used when `--language` is omitted. diff --git a/scripts/nemo/nemotron-3.5-asr-streaming-0.6b/export_onnx.py b/scripts/nemo/nemotron-3.5-asr-streaming-0.6b/export_onnx.py new file mode 100644 index 0000000000..2a4083940f --- /dev/null +++ b/scripts/nemo/nemotron-3.5-asr-streaming-0.6b/export_onnx.py @@ -0,0 +1,544 @@ +#!/usr/bin/env python3 +# Copyright 2026 Julian Pscheid +import inspect +import json +import os +from pathlib import Path +from typing import Any, Dict, Optional + +import nemo.collections.asr as nemo_asr +import onnx +import torch +from onnxruntime.quantization import QuantType, quantize_dynamic + +ENCODER_INPUT_NAMES = [ + "audio_signal", + "length", + "cache_last_channel", + "cache_last_time", + "cache_last_channel_len", + "prompt_index", +] + +ENCODER_OUTPUT_NAMES = [ + "outputs", + "encoded_lengths", + "cache_last_channel_next", + "cache_last_time_next", + "cache_last_channel_next_len", +] + +FORWARD_FOR_EXPORT_ARGS = [ + "audio_signal", + "length", + "cache_last_channel", + "cache_last_time", + "cache_last_channel_len", +] + + +def add_meta_data(filename: str, meta_data: Dict[str, str]): + """Add meta data to an ONNX model. It is changed in-place.""" + model = onnx.load(filename) + + while len(model.metadata_props): + model.metadata_props.pop() + + for key, value in meta_data.items(): + meta = model.metadata_props.add() + meta.key = key + meta.value = str(value) + + external_filename = filename.split(".onnx")[0] + # onnx.save refuses to overwrite an existing external-data file; the + # prompted-encoder export already wrote one, so remove it first. + Path(external_filename + ".data").unlink(missing_ok=True) + onnx.save( + model, + filename, + save_as_external_data=True, + all_tensors_to_one_file=True, + location=external_filename + ".data", + ) + + +def _to_plain_container(obj: Any) -> Any: + try: + from omegaconf import DictConfig, ListConfig, OmegaConf + + if isinstance(obj, (DictConfig, ListConfig)): + return OmegaConf.to_container(obj, resolve=True) + except ImportError: + pass + + return obj + + +def _normalize_prompt_dictionary(obj: Any) -> Dict[str, int]: + obj = _to_plain_container(obj) + if not isinstance(obj, dict): + raise TypeError(type(obj)) + + return {str(k): int(v) for k, v in obj.items()} + + +def _get_config_value(obj: Any, key: str) -> Any: + if obj is None: + return None + + try: + return getattr(obj, key) + except (AttributeError, KeyError): + pass + + obj = _to_plain_container(obj) + if isinstance(obj, dict): + return obj.get(key) + + return None + + +def get_prompt_dictionary(asr_model) -> Dict[str, int]: + """Return the model's language prompt dictionary from NeMo artifacts.""" + cfg = getattr(asr_model, "cfg", None) + model_defaults = _get_config_value(cfg, "model_defaults") + if model_defaults is None: + raise RuntimeError("Could not find cfg.model_defaults in the NeMo model") + + prompt_dictionary = _get_config_value(model_defaults, "prompt_dictionary") + if prompt_dictionary is None: + raise RuntimeError( + "Could not find cfg.model_defaults.prompt_dictionary in the NeMo model" + ) + + try: + ans = _normalize_prompt_dictionary(prompt_dictionary) + except (TypeError, ValueError) as e: + raise RuntimeError( + "cfg.model_defaults.prompt_dictionary must map language strings " + "to integer prompt ids" + ) from e + + num_prompts = int(asr_model.num_prompts) + for language, prompt_id in ans.items(): + if not 0 <= prompt_id < num_prompts: + raise ValueError( + "cfg.model_defaults.prompt_dictionary has out-of-range " + f"prompt id for '{language}': {prompt_id}; expected " + f"0 <= id < {num_prompts}" + ) + + auto_prompt_id = ans.get("auto") + if auto_prompt_id != 101: + raise ValueError(f"Expected auto prompt id 101, got {auto_prompt_id}") + + # The dictionary may use locale-style keys such as en-US or ja-JP; the + # runtime derives base-code aliases, so accept either form here. + for language in ["en", "ja"]: + if not any(k == language or k.startswith(f"{language}-") for k in ans): + raise RuntimeError( + "cfg.model_defaults.prompt_dictionary is missing " + f"'{language}'" + ) + + return ans + + +def _find_sentencepiece_processor(obj: Any, max_depth: int = 5) -> Optional[Any]: + seen = set() + + def is_sentencepiece_processor(value: Any) -> bool: + return callable(getattr(value, "get_piece_size", None)) and callable( + getattr(value, "id_to_piece", None) + ) + + def visit(value: Any, depth: int) -> Optional[Any]: + if value is None or depth > max_depth: + return None + + if is_sentencepiece_processor(value): + return value + + obj_id = id(value) + if obj_id in seen: + return None + seen.add(obj_id) + + for name in ["tokenizer", "sp_model", "model", "processor"]: + if hasattr(value, name): + found = visit(getattr(value, name), depth + 1) + if found is not None: + return found + + if isinstance(value, dict): + for v in value.values(): + found = visit(v, depth + 1) + if found is not None: + return found + + return None + + return visit(obj, 0) + + +def save_tokens(asr_model, filename: str = "tokens.txt") -> int: + sp = _find_sentencepiece_processor(getattr(asr_model, "tokenizer", None)) + if sp is None: + raise RuntimeError("Could not find the SentencePiece tokenizer in the model") + + vocab_size = sp.get_piece_size() + with open(filename, "w", encoding="utf-8") as f: + for i in range(vocab_size): + f.write(f"{sp.id_to_piece(i)} {i}\n") + f.write(f" {vocab_size}\n") + + print(f"Saved {filename}") + return vocab_size + + +def assert_forward_for_export_signature(encoder): + if not hasattr(encoder, "forward_for_export"): + raise RuntimeError("Expected encoder.forward_for_export for ONNX export") + + signature = inspect.signature(encoder.forward_for_export) + missing = [ + name for name in FORWARD_FOR_EXPORT_ARGS if name not in signature.parameters + ] + if missing: + raise RuntimeError( + "encoder.forward_for_export is missing expected argument(s): " + f"{missing}. Signature: {signature}" + ) + + +class PromptedStreamingEncoder(torch.nn.Module): + def __init__(self, asr_model): + super().__init__() + + for attr in ["encoder", "prompt_kernel", "num_prompts"]: + if not hasattr(asr_model, attr): + raise RuntimeError( + "Expected a prompt-conditioned NeMo model with " + f"'{attr}'" + ) + + self.encoder = asr_model.encoder + assert_forward_for_export_signature(self.encoder) + + self.prompt_kernel = asr_model.prompt_kernel + self.num_prompts = int(asr_model.num_prompts) + + def forward( + self, + audio_signal, + length, + cache_last_channel, + cache_last_time, + cache_last_channel_len, + prompt_index, + ): + encoded, encoded_len, channel_next, time_next, channel_len_next = ( + self.encoder.forward_for_export( + audio_signal=audio_signal, + length=length, + cache_last_channel=cache_last_channel, + cache_last_time=cache_last_time, + cache_last_channel_len=cache_last_channel_len, + ) + ) + + # Mirror NeMo PromptStreamingMixin._apply_prompt_to_encoded(), but make + # the prompt id a real ONNX input instead of a Python inference setting. + out_dtype = encoded.dtype + encoded = encoded.transpose(1, 2) # (B, D, T) -> (B, T, D) + batch_size, time_steps, _ = encoded.shape + + prompt = torch.zeros( + batch_size, + time_steps, + self.num_prompts, + dtype=encoded.dtype, + device=encoded.device, + ) + prompt.scatter_( + 2, + prompt_index.view(batch_size, 1, 1).expand(-1, time_steps, -1), + 1.0, + ) + + encoded = self.prompt_kernel(torch.cat([encoded, prompt], dim=-1)).to( + out_dtype + ) + encoded = encoded.transpose(1, 2) # (B, T, D) -> (B, D, T) + + return encoded, encoded_len, channel_next, time_next, channel_len_next + + +def remove_export_scratch_files(): + patterns = [ + "Constant_*_attr__value", + "onnx__MatMul_*", + "layers.*.conv*", + "pre_encode.conv.*.weight", + "encoder.export.onnx*", + ] + for pattern in patterns: + for p in Path(".").glob(pattern): + p.unlink() + + +def assert_encoder_graph(filename: str): + model = onnx.load(filename, load_external_data=False) + + input_names = [i.name for i in model.graph.input] + if input_names != ENCODER_INPUT_NAMES: + raise RuntimeError( + f"{filename}: expected encoder inputs {ENCODER_INPUT_NAMES}, " + f"got {input_names}" + ) + + output_names = [o.name for o in model.graph.output] + if output_names != ENCODER_OUTPUT_NAMES: + raise RuntimeError( + f"{filename}: expected encoder outputs {ENCODER_OUTPUT_NAMES}, " + f"got {output_names}" + ) + + +def _module_device_and_dtype(module): + try: + p = next(module.parameters()) + return p.device, p.dtype + except StopIteration: + return torch.device("cpu"), torch.float32 + + +def export_prompted_encoder( + asr_model, + window_size: int, + cache_last_channel_dim1: int, + cache_last_channel_dim2: int, + cache_last_channel_dim3: int, + cache_last_time_dim1: int, + cache_last_time_dim2: int, + cache_last_time_dim3: int, + auto_prompt_id: int, +): + device, dtype = _module_device_and_dtype(asr_model.encoder) + + audio_signal = torch.zeros( + 1, 128, window_size, dtype=dtype, device=device + ) + length = torch.full((1,), window_size, dtype=torch.int64, device=device) + cache_last_channel = torch.zeros( + 1, + cache_last_channel_dim1, + cache_last_channel_dim2, + cache_last_channel_dim3, + dtype=dtype, + device=device, + ) + cache_last_time = torch.zeros( + 1, + cache_last_time_dim1, + cache_last_time_dim2, + cache_last_time_dim3, + dtype=dtype, + device=device, + ) + cache_last_channel_len = torch.zeros(1, dtype=torch.int64, device=device) + prompt_index = torch.full((1,), auto_prompt_id, dtype=torch.int64, device=device) + + encoder = PromptedStreamingEncoder(asr_model).eval() + + export_kwargs = {} + export_signature = inspect.signature(torch.onnx.export) + if "dynamo" in export_signature.parameters: + export_kwargs["dynamo"] = False + if "external_data" in export_signature.parameters: + export_kwargs["external_data"] = True + + torch.onnx.export( + encoder, + ( + audio_signal, + length, + cache_last_channel, + cache_last_time, + cache_last_channel_len, + prompt_index, + ), + "encoder.export.onnx", + input_names=ENCODER_INPUT_NAMES, + output_names=ENCODER_OUTPUT_NAMES, + opset_version=17, + dynamic_axes={ + "audio_signal": {0: "batch", 2: "time"}, + "length": {0: "batch"}, + "cache_last_channel": {0: "batch", 2: "cache_channel_time"}, + "cache_last_time": {0: "batch", 3: "cache_time_width"}, + "cache_last_channel_len": {0: "batch"}, + "prompt_index": {0: "batch"}, + "outputs": {0: "batch", 2: "time"}, + "encoded_lengths": {0: "batch"}, + "cache_last_channel_next": {0: "batch", 2: "cache_channel_time"}, + "cache_last_time_next": {0: "batch", 3: "cache_time_width"}, + "cache_last_channel_next_len": {0: "batch"}, + }, + **export_kwargs, + ) + + model = onnx.load("encoder.export.onnx", load_external_data=True) + onnx.save_model( + model, + "encoder.onnx", + save_as_external_data=True, + all_tensors_to_one_file=True, + location="encoder.data", + size_threshold=0, + ) + assert_encoder_graph("encoder.onnx") + for p in Path(".").glob("encoder.export.onnx*"): + p.unlink() + + +@torch.no_grad() +def main(): + model_name = "nvidia/nemotron-3.5-asr-streaming-0.6b" + + asr_model = nemo_asr.models.ASRModel.from_pretrained(model_name=model_name) + + vocab_size = save_tokens(asr_model) + if vocab_size != asr_model.decoder.vocab_size: + raise ValueError( + f"SentencePiece vocab size {vocab_size} != decoder vocab size " + f"{asr_model.decoder.vocab_size}" + ) + + prompt_dictionary = get_prompt_dictionary(asr_model) + auto_prompt_id = prompt_dictionary["auto"] + if auto_prompt_id != 101: + raise ValueError(f"Expected auto prompt id 101, got {auto_prompt_id}") + + asr_model.eval() + + assert asr_model.encoder.streaming_cfg is not None + print("streaming_cfg", asr_model.encoder.streaming_cfg) + print("prompt_dictionary", prompt_dictionary) + + chunk_size_ms_list = [80, 160, 560, 1120] + for ms in chunk_size_ms_list: + chunk_size = ms // 80 - 1 + print("chunk_size", chunk_size) + asr_model.encoder.set_default_att_context_size([70, chunk_size]) + + print("streaming_cfg", asr_model.encoder.streaming_cfg) + print("att_context_size", asr_model.encoder.att_context_size) + print( + "pre_encode_cache_size", + asr_model.encoder.streaming_cfg.pre_encode_cache_size, + ) + + if isinstance(asr_model.encoder.streaming_cfg.pre_encode_cache_size, list): + pre_encode_cache_size = ( + asr_model.encoder.streaming_cfg.pre_encode_cache_size[1] + ) + else: + pre_encode_cache_size = ( + asr_model.encoder.streaming_cfg.pre_encode_cache_size + ) + + if isinstance(asr_model.encoder.streaming_cfg.chunk_size, list): + chunk_size = asr_model.encoder.streaming_cfg.chunk_size[1] + else: + chunk_size = asr_model.encoder.streaming_cfg.chunk_size + + window_size = chunk_size + pre_encode_cache_size + + print("chunk_size", chunk_size) + print("pre_encode_cache_size", pre_encode_cache_size) + print("window_size", window_size) + + chunk_shift = chunk_size + + # cache_last_channel: (batch_size, dim1, dim2, dim3) + cache_last_channel_dim1 = len(asr_model.encoder.layers) + cache_last_channel_dim2 = ( + asr_model.encoder.streaming_cfg.last_channel_cache_size + ) + cache_last_channel_dim3 = asr_model.encoder.d_model + + # cache_last_time: (batch_size, dim1, dim2, dim3) + cache_last_time_dim1 = len(asr_model.encoder.layers) + cache_last_time_dim2 = asr_model.encoder.d_model + cache_last_time_dim3 = asr_model.encoder.conv_context_size[0] + + asr_model.set_export_config({"cache_support": True}) + + export_prompted_encoder( + asr_model=asr_model, + window_size=window_size, + cache_last_channel_dim1=cache_last_channel_dim1, + cache_last_channel_dim2=cache_last_channel_dim2, + cache_last_channel_dim3=cache_last_channel_dim3, + cache_last_time_dim1=cache_last_time_dim1, + cache_last_time_dim2=cache_last_time_dim2, + cache_last_time_dim3=cache_last_time_dim3, + auto_prompt_id=auto_prompt_id, + ) + asr_model.decoder.export("decoder.onnx") + asr_model.joint.export("joiner.onnx") + + normalize_type = asr_model.cfg.preprocessor.normalize + if normalize_type == "NA": + normalize_type = "" + + meta_data = { + "vocab_size": asr_model.decoder.vocab_size, + "window_size": window_size, + "chunk_size_ms": ms, + "chunk_shift": chunk_shift, + "normalize_type": normalize_type, + "cache_last_channel_dim1": cache_last_channel_dim1, + "cache_last_channel_dim2": cache_last_channel_dim2, + "cache_last_channel_dim3": cache_last_channel_dim3, + "cache_last_time_dim1": cache_last_time_dim1, + "cache_last_time_dim2": cache_last_time_dim2, + "cache_last_time_dim3": cache_last_time_dim3, + "pred_rnn_layers": asr_model.decoder.pred_rnn_layers, + "pred_hidden": asr_model.decoder.pred_hidden, + "subsampling_factor": 8, + "feat_dim": 128, + "model_type": type(asr_model).__name__, + "version": "1", + "model_author": "NeMo", + "url": f"https://huggingface.co/{model_name}", + "comment": "Only the transducer branch is exported", + "prompt_dictionary": json.dumps(prompt_dictionary, sort_keys=True), + "auto_prompt_id": auto_prompt_id, + } + print("meta_data", meta_data) + add_meta_data("encoder.onnx", meta_data) + assert_encoder_graph("encoder.onnx") + + for m in ["encoder", "decoder", "joiner"]: + quantize_dynamic( + model_input=f"{m}.onnx", + model_output=f"{m}.int8.onnx", + weight_type=QuantType.QUInt8, + ) + assert_encoder_graph("encoder.int8.onnx") + + Path(str(ms)).mkdir(exist_ok=True) + for suffix in ["onnx", "data"]: + for p in Path(".").glob(f"*.{suffix}"): + p.rename(Path(str(ms)) / p.name) + + print(meta_data) + print(f"Saved exported models to {ms}") + remove_export_scratch_files() + os.system(f"ls -lh {ms}") + + +if __name__ == "__main__": + main() diff --git a/sherpa-onnx/csrc/online-recognizer-transducer-nemo-impl.h b/sherpa-onnx/csrc/online-recognizer-transducer-nemo-impl.h index e2132a38f8..2cfe4148ee 100644 --- a/sherpa-onnx/csrc/online-recognizer-transducer-nemo-impl.h +++ b/sherpa-onnx/csrc/online-recognizer-transducer-nemo-impl.h @@ -7,12 +7,14 @@ #define SHERPA_ONNX_CSRC_ONLINE_RECOGNIZER_TRANSDUCER_NEMO_IMPL_H_ #include +#include #include #include #include #include // NOLINT #include #include +#include #include #include @@ -100,9 +102,17 @@ class OnlineRecognizerTransducerNeMoImpl : public OnlineRecognizerImpl { // TODO(fangjun): Remember to change these constants if needed int32_t frame_shift_ms = 10; int32_t subsampling_factor = model_->SubsamplingFactor(); - auto r = Convert(s->GetResult(), symbol_table_, frame_shift_ms, - subsampling_factor, s->GetCurrentSegment(), - s->GetNumFramesSinceStart()); + const auto &decoder_result = s->GetResult(); + bool has_language_tag = !language_tag_token_ids_.empty() && + ContainsLanguageTag(decoder_result); + auto r = has_language_tag + ? Convert(FilterLanguageTags(decoder_result), symbol_table_, + frame_shift_ms, subsampling_factor, + s->GetCurrentSegment(), s->GetNumFramesSinceStart()) + : Convert(decoder_result, symbol_table_, frame_shift_ms, + subsampling_factor, s->GetCurrentSegment(), + s->GetNumFramesSinceStart()); + r.text = ApplyInverseTextNormalization(std::move(r.text)); r.text = ApplyHomophoneReplacer(std::move(r.text)); return r; @@ -180,7 +190,9 @@ class OnlineRecognizerTransducerNeMoImpl : public OnlineRecognizerImpl { auto states = model_->StackStates(std::move(encoder_states)); int32_t num_states = states.size(); // num_states = 3 - auto t = model_->RunEncoder(std::move(x), std::move(states)); + auto language_prompt_ids = GetLanguagePromptIds(ss, n); + auto t = model_->RunEncoder(std::move(x), std::move(states), + language_prompt_ids); // t[0] encoder_out, float tensor, (batch_size, dim, T) // t[1] next states @@ -210,6 +222,127 @@ class OnlineRecognizerTransducerNeMoImpl : public OnlineRecognizerImpl { } private: + static bool IsLanguageTagToken(const std::string &sym) { + if (sym.size() < 4 || sym.front() != '<' || sym.back() != '>') { + return false; + } + + size_t i = 1; + int32_t num_lowercase = 0; + while (i + 1 < sym.size() && num_lowercase != 3 && + std::islower(static_cast(sym[i]))) { + ++i; + ++num_lowercase; + } + + if (num_lowercase < 2) { + return false; + } + + if (i == sym.size() - 1) { + return true; + } + + if (sym[i] != '-' || i + 3 != sym.size() - 1) { + return false; + } + + return std::isupper(static_cast(sym[i + 1])) && + std::isupper(static_cast(sym[i + 2])); + } + + void InitLanguageTagTokenIds() { + if (!model_->IsMultilingual()) { + return; + } + + for (int32_t i = 0; i != symbol_table_.NumSymbols(); ++i) { + if (symbol_table_.Contains(i) && IsLanguageTagToken(symbol_table_[i])) { + language_tag_token_ids_.insert(i); + } + } + } + + bool ContainsLanguageTag(const OnlineTransducerDecoderResult &src) const { + for (auto token : src.tokens) { + if (language_tag_token_ids_.count(token)) { + return true; + } + } + + return false; + } + + OnlineTransducerDecoderResult FilterLanguageTags( + const OnlineTransducerDecoderResult &src) const { + OnlineTransducerDecoderResult ans; + ans.frame_offset = src.frame_offset; + ans.num_trailing_blanks = src.num_trailing_blanks; + ans.tokens.reserve(src.tokens.size()); + ans.timestamps.reserve(src.timestamps.size()); + + bool filter_ys_probs = src.ys_probs.size() == src.tokens.size(); + bool filter_lm_probs = src.lm_probs.size() == src.tokens.size(); + bool filter_context_scores = + src.context_scores.size() == src.tokens.size(); + + if (filter_ys_probs) { + ans.ys_probs.reserve(src.ys_probs.size()); + } else { + ans.ys_probs = src.ys_probs; + } + + if (filter_lm_probs) { + ans.lm_probs.reserve(src.lm_probs.size()); + } else { + ans.lm_probs = src.lm_probs; + } + + if (filter_context_scores) { + ans.context_scores.reserve(src.context_scores.size()); + } else { + ans.context_scores = src.context_scores; + } + + for (size_t i = 0; i != src.tokens.size(); ++i) { + if (language_tag_token_ids_.count(src.tokens[i])) { + continue; + } + + ans.tokens.push_back(src.tokens[i]); + if (i < src.timestamps.size()) { + ans.timestamps.push_back(src.timestamps[i]); + } + + if (filter_ys_probs) { + ans.ys_probs.push_back(src.ys_probs[i]); + } + if (filter_lm_probs) { + ans.lm_probs.push_back(src.lm_probs[i]); + } + if (filter_context_scores) { + ans.context_scores.push_back(src.context_scores[i]); + } + } + + return ans; + } + + std::vector GetLanguagePromptIds(OnlineStream **ss, + int32_t n) const { + std::vector ans; + if (!model_->IsMultilingual()) { + return ans; + } + + ans.reserve(n); + for (int32_t i = 0; i != n; ++i) { + ans.push_back(model_->GetLanguagePromptId(ss[i]->GetOption("language"))); + } + + return ans; + } + void PostInit() { config_.feat_config.feature_dim = model_->FeatureDim(); @@ -240,12 +373,15 @@ class OnlineRecognizerTransducerNeMoImpl : public OnlineRecognizerImpl { symbol_table_.NumSymbols(), vocab_size); SHERPA_ONNX_EXIT(-1); } + + InitLanguageTagTokenIds(); } private: OnlineRecognizerConfig config_; SymbolTable symbol_table_; std::unique_ptr model_; + std::unordered_set language_tag_token_ids_; std::unique_ptr decoder_; Endpoint endpoint_; }; diff --git a/sherpa-onnx/csrc/online-transducer-greedy-search-nemo-decoder.cc b/sherpa-onnx/csrc/online-transducer-greedy-search-nemo-decoder.cc index 855e76c253..1a2eeddf49 100644 --- a/sherpa-onnx/csrc/online-transducer-greedy-search-nemo-decoder.cc +++ b/sherpa-onnx/csrc/online-transducer-greedy-search-nemo-decoder.cc @@ -29,6 +29,17 @@ static Ort::Value BuildDecoderInput(int32_t token, OrtAllocator *allocator) { return decoder_input; } +static std::vector BuildStateViews( + std::vector *states) { + std::vector ans; + ans.reserve(states->size()); + for (auto &v : *states) { + ans.push_back(View(&v)); + } + + return ans; +} + static void DecodeOne(const float *encoder_out, int32_t num_rows, int32_t num_cols, OnlineTransducerNeMoModel *model, float blank_penalty, OnlineStream *s) { @@ -48,10 +59,7 @@ static void DecodeOne(const float *encoder_out, int32_t num_rows, std::vector &last_decoder_states = s->GetNeMoDecoderStates(); std::vector tmp_decoder_states; - tmp_decoder_states.reserve(last_decoder_states.size()); - for (auto &v : last_decoder_states) { - tmp_decoder_states.push_back(View(&v)); - } + tmp_decoder_states = BuildStateViews(&last_decoder_states); // decoder_output_pair.second returns the next decoder state std::pair> decoder_output_pair = diff --git a/sherpa-onnx/csrc/online-transducer-nemo-model.cc b/sherpa-onnx/csrc/online-transducer-nemo-model.cc index e13f2c8a5e..9a543778ac 100644 --- a/sherpa-onnx/csrc/online-transducer-nemo-model.cc +++ b/sherpa-onnx/csrc/online-transducer-nemo-model.cc @@ -9,10 +9,12 @@ #include #include #include +#include #include #include #include #include +#include #include #include @@ -25,6 +27,7 @@ #include "rawfile/raw_file_manager.h" #endif +#include "nlohmann/json.hpp" #include "sherpa-onnx/csrc/cat.h" #include "sherpa-onnx/csrc/file-utils.h" #include "sherpa-onnx/csrc/macros.h" @@ -37,6 +40,105 @@ namespace sherpa_onnx { +namespace { + +constexpr int64_t kDefaultAutoPromptId = 101; + +std::string StripQuotes(std::string s) { + s = Trim(s); + if (s.size() >= 2 && + ((s.front() == '"' && s.back() == '"') || + (s.front() == '\'' && s.back() == '\''))) { + return s.substr(1, s.size() - 2); + } + return s; +} + +std::string NormalizeLanguage(std::string s) { + s = StripQuotes(std::move(s)); + s = Trim(s); + if (s.size() >= 2 && s.front() == '<' && s.back() == '>') { + s = s.substr(1, s.size() - 2); + } + + std::replace(s.begin(), s.end(), '_', '-'); + ToLowerCase(&s); + return s; +} + +void AddLanguagePromptId( + const std::string &language, int64_t prompt_id, + std::unordered_map *language_prompt_ids, + std::vector> *ordered_prompt_ids) { + auto normalized = NormalizeLanguage(language); + if (normalized.empty()) { + return; + } + + if (language_prompt_ids->emplace(normalized, prompt_id).second) { + ordered_prompt_ids->push_back({normalized, prompt_id}); + } +} + +void AddBaseLanguageAliases( + std::unordered_map *language_prompt_ids, + const std::vector> &ordered_prompt_ids) { + for (const auto &p : ordered_prompt_ids) { + auto pos = p.first.find('-'); + if (pos == std::string::npos || pos == 0) { + continue; + } + + auto base = p.first.substr(0, pos); + language_prompt_ids->emplace(std::move(base), p.second); + } +} + +// The metadata value is a JSON object mapping language strings to integer +// prompt ids, e.g. {"auto": 101, "en-US": 0, "ja-JP": 1}. +void ParseLanguagePromptDictionary( + const std::string &value, + std::unordered_map *language_prompt_ids, + std::vector> *ordered_prompt_ids) { + auto j = nlohmann::json::parse(value, nullptr, /*allow_exceptions*/ false); + if (j.is_discarded() || !j.is_object()) { + return; + } + + for (const auto &item : j.items()) { + if (!item.value().is_number_integer()) { + continue; + } + + int64_t prompt_id = item.value().get(); + if (prompt_id < 0) { + continue; + } + + AddLanguagePromptId(item.key(), prompt_id, language_prompt_ids, + ordered_prompt_ids); + } +} + +bool ReadPromptIdFromMetadata(const Ort::ModelMetadata &meta_data, + OrtAllocator *allocator, const char *key, + int64_t *prompt_id) { + auto value = LookupCustomModelMetaData(meta_data, key, allocator); + if (value.empty()) { + return false; + } + + int64_t id = -1; + if (!ConvertStringToInteger(Trim(value), &id) || id < 0) { + return false; + } + + *prompt_id = id; + return true; +} + +} // namespace + class OnlineTransducerNeMoModel::Impl { public: explicit Impl(const OnlineModelConfig &config) @@ -73,14 +175,13 @@ class OnlineTransducerNeMoModel::Impl { InitDecoder(buf.data(), buf.size()); } - { - auto buf = ReadFile(mgr, config.transducer.joiner); - InitJoiner(buf.data(), buf.size()); - } + auto buf = ReadFile(mgr, config.transducer.joiner); + InitJoiner(buf.data(), buf.size()); } - std::vector RunEncoder(Ort::Value features, - std::vector states) { + std::vector RunEncoder( + Ort::Value features, std::vector states, + const std::vector &language_prompt_ids) { Ort::Value &cache_last_channel = states[0]; Ort::Value &cache_last_time = states[1]; Ort::Value &cache_last_channel_len = states[2]; @@ -99,9 +200,28 @@ class OnlineTransducerNeMoModel::Impl { // (B, T, C) -> (B, C, T) features = Transpose12(allocator_, &features); - std::array inputs = { - std::move(features), View(&length), std::move(cache_last_channel), - std::move(cache_last_time), std::move(cache_last_channel_len)}; + std::vector inputs; + inputs.reserve(is_multilingual_ ? 6 : 5); + inputs.push_back(std::move(features)); + inputs.push_back(View(&length)); + inputs.push_back(std::move(cache_last_channel)); + inputs.push_back(std::move(cache_last_time)); + inputs.push_back(std::move(cache_last_channel_len)); + + Ort::Value prompt_id_tensor{nullptr}; + if (is_multilingual_) { + std::array prompt_id_shape{batch_size}; + prompt_id_tensor = Ort::Value::CreateTensor( + allocator_, prompt_id_shape.data(), prompt_id_shape.size()); + int64_t *p_prompt_id = prompt_id_tensor.GetTensorMutableData(); + if (static_cast(language_prompt_ids.size()) == batch_size) { + std::copy(language_prompt_ids.begin(), language_prompt_ids.end(), + p_prompt_id); + } else { + std::fill(p_prompt_id, p_prompt_id + batch_size, default_prompt_id_); + } + inputs.push_back(std::move(prompt_id_tensor)); + } auto out = encoder_sess_->Run( {}, encoder_input_names_ptr_.data(), inputs.data(), inputs.size(), @@ -199,6 +319,35 @@ class OnlineTransducerNeMoModel::Impl { int32_t FeatureDim() const { return feat_dim_; } + bool IsMultilingual() const { return is_multilingual_; } + + int64_t GetLanguagePromptId(const std::string &language) const { + if (!is_multilingual_ || language.empty()) { + return default_prompt_id_; + } + + auto it = language_prompt_ids_.find(language); + if (it != language_prompt_ids_.end()) { + return it->second; + } + + auto normalized = NormalizeLanguage(language); + if (normalized.empty() || normalized == "auto") { + return default_prompt_id_; + } + + it = language_prompt_ids_.find(normalized); + if (it != language_prompt_ids_.end()) { + return it->second; + } + + SHERPA_ONNX_LOGE( + "Unsupported language '%s' for multilingual NeMo transducer; using " + "auto", + language.c_str()); + return default_prompt_id_; + } + int32_t VocabSize() const { return vocab_size_; } OrtAllocator *Allocator() { return allocator_; } @@ -307,6 +456,10 @@ class OnlineTransducerNeMoModel::Impl { GetOutputNames(encoder_sess_.get(), &encoder_output_names_, &encoder_output_names_ptr_); + is_multilingual_ = + std::find(encoder_input_names_.begin(), encoder_input_names_.end(), + "prompt_index") != encoder_input_names_.end(); + feat_dim_ = encoder_sess_->GetInputTypeInfo(0) .GetTensorTypeAndShapeInfo() .GetShape()[1]; @@ -318,6 +471,7 @@ class OnlineTransducerNeMoModel::Impl { os << "---encoder---\n"; PrintModelMetadata(os, meta_data); os << "feat_dim: " << feat_dim_ << "\n"; + os << "is_multilingual: " << (is_multilingual_ ? "1" : "0") << "\n"; #if __OHOS__ SHERPA_ONNX_LOGE("%{public}s", os.str().c_str()); #else @@ -355,11 +509,52 @@ class OnlineTransducerNeMoModel::Impl { normalize_type_ = ""; } + if (is_multilingual_) { + InitLanguagePromptIds(meta_data, allocator); + } + InitEncoderStates(); } + void InitLanguagePromptIds(const Ort::ModelMetadata &meta_data, + OrtAllocator *allocator) { + auto dict = LookupCustomModelMetaData(meta_data, "prompt_dictionary", + allocator); + if (!dict.empty()) { + ParseLanguagePromptDictionary(dict, &language_prompt_ids_, + &ordered_prompt_ids_); + } + + ReadPromptIdFromMetadata(meta_data, allocator, "auto_prompt_id", + &default_prompt_id_); + + auto it = language_prompt_ids_.find("auto"); + if (it != language_prompt_ids_.end()) { + default_prompt_id_ = it->second; + } else { + AddLanguagePromptId("auto", default_prompt_id_, &language_prompt_ids_, + &ordered_prompt_ids_); + } + + bool has_non_auto_prompt = false; + for (const auto &p : ordered_prompt_ids_) { + if (p.first != "auto") { + has_non_auto_prompt = true; + break; + } + } + if (!has_non_auto_prompt) { + SHERPA_ONNX_LOGE( + "The encoder declares prompt_index, but usable prompt_dictionary " + "metadata is missing; all languages will fall back to auto."); + } + + AddBaseLanguageAliases(&language_prompt_ids_, ordered_prompt_ids_); + } + void InitEncoderStates() { - std::array cache_last_channel_shape{1, cache_last_channel_dim1_, + std::array cache_last_channel_shape{1, + cache_last_channel_dim1_, cache_last_channel_dim2_, cache_last_channel_dim3_}; @@ -471,6 +666,10 @@ class OnlineTransducerNeMoModel::Impl { int32_t vocab_size_ = 0; int32_t subsampling_factor_ = 8; int32_t feat_dim_ = 80; + bool is_multilingual_ = false; + int64_t default_prompt_id_ = kDefaultAutoPromptId; + std::unordered_map language_prompt_ids_; + std::vector> ordered_prompt_ids_; std::string normalize_type_; int32_t pred_rnn_layers_ = -1; int32_t pred_hidden_ = -1; @@ -505,8 +704,10 @@ OnlineTransducerNeMoModel::OnlineTransducerNeMoModel( OnlineTransducerNeMoModel::~OnlineTransducerNeMoModel() = default; std::vector OnlineTransducerNeMoModel::RunEncoder( - Ort::Value features, std::vector states) const { - return impl_->RunEncoder(std::move(features), std::move(states)); + Ort::Value features, std::vector states, + const std::vector &language_prompt_ids) const { + return impl_->RunEncoder(std::move(features), std::move(states), + language_prompt_ids); } std::pair> @@ -545,6 +746,15 @@ int32_t OnlineTransducerNeMoModel::FeatureDim() const { return impl_->FeatureDim(); } +bool OnlineTransducerNeMoModel::IsMultilingual() const { + return impl_->IsMultilingual(); +} + +int64_t OnlineTransducerNeMoModel::GetLanguagePromptId( + const std::string &language) const { + return impl_->GetLanguagePromptId(language); +} + OrtAllocator *OnlineTransducerNeMoModel::Allocator() const { return impl_->Allocator(); } diff --git a/sherpa-onnx/csrc/online-transducer-nemo-model.h b/sherpa-onnx/csrc/online-transducer-nemo-model.h index ff38413323..183364bbf2 100644 --- a/sherpa-onnx/csrc/online-transducer-nemo-model.h +++ b/sherpa-onnx/csrc/online-transducer-nemo-model.h @@ -6,6 +6,7 @@ #ifndef SHERPA_ONNX_CSRC_ONLINE_TRANSDUCER_NEMO_MODEL_H_ #define SHERPA_ONNX_CSRC_ONLINE_TRANSDUCER_NEMO_MODEL_H_ +#include #include #include #include @@ -52,7 +53,8 @@ class OnlineTransducerNeMoModel { * - ans[1:]: contains next states */ std::vector RunEncoder( - Ort::Value features, std::vector states) const; // NOLINT + Ort::Value features, std::vector states, // NOLINT + const std::vector &language_prompt_ids = {}) const; /** Run the decoder network. * @@ -101,6 +103,10 @@ class OnlineTransducerNeMoModel { int32_t FeatureDim() const; + bool IsMultilingual() const; + + int64_t GetLanguagePromptId(const std::string &language) const; + /** Return an allocator for allocating memory */ OrtAllocator *Allocator() const; diff --git a/sherpa-onnx/csrc/sherpa-onnx.cc b/sherpa-onnx/csrc/sherpa-onnx.cc index 4b116472dd..e3c2631402 100644 --- a/sherpa-onnx/csrc/sherpa-onnx.cc +++ b/sherpa-onnx/csrc/sherpa-onnx.cc @@ -81,8 +81,12 @@ for a list of pre-trained models to download. sherpa_onnx::ParseOptions po(kUsageMessage); sherpa_onnx::OnlineRecognizerConfig config; + std::string language; config.Register(&po); + po.Register("language", &language, + "Per-stream language hint for prompt-conditioned multilingual " + "models, e.g., en, fr, ja, or auto. Empty means auto."); po.Read(argc, argv); if (po.NumArgs() < 1) { @@ -124,10 +128,13 @@ for a list of pre-trained models to download. const float duration = samples.size() / static_cast(sampling_rate); auto s = recognizer.CreateStream(); + if (!language.empty()) { + s->SetOption("language", language); + } - // std::vector left_paddings(static_cast(0.3 * sampling_rate)); - // s->AcceptWaveform(sampling_rate, left_paddings.data(), - // left_paddings.size()); + std::vector left_paddings(static_cast(0.3 * sampling_rate)); + s->AcceptWaveform(sampling_rate, left_paddings.data(), + left_paddings.size()); s->AcceptWaveform(sampling_rate, samples.data(), samples.size()); diff --git a/sherpa-onnx/python/tests/test_online_recognizer.py b/sherpa-onnx/python/tests/test_online_recognizer.py index 5319f41eaa..7c7bb80e73 100755 --- a/sherpa-onnx/python/tests/test_online_recognizer.py +++ b/sherpa-onnx/python/tests/test_online_recognizer.py @@ -6,6 +6,7 @@ # # ctest --verbose -R test_online_recognizer_py +import re import unittest import wave from pathlib import Path @@ -143,6 +144,97 @@ def test_transducer_multiple_files(self): print(f"{wave_filename}\n{result}") print("-" * 10) + def test_nemotron_streaming_english(self): + m = "sherpa-onnx-nemotron-speech-streaming-en-0.6b-560ms-int8-2026-04-25" + encoder = f"{d}/{m}/encoder.int8.onnx" + decoder = f"{d}/{m}/decoder.int8.onnx" + joiner = f"{d}/{m}/joiner.int8.onnx" + tokens = f"{d}/{m}/tokens.txt" + wave0 = f"{d}/{m}/test_wavs/0.wav" + + if not Path(encoder).is_file(): + print("skipping test_nemotron_streaming_english()") + return + + recognizer = sherpa_onnx.OnlineRecognizer.from_transducer( + encoder=encoder, + decoder=decoder, + joiner=joiner, + tokens=tokens, + num_threads=1, + provider="cpu", + ) + s = recognizer.create_stream() + samples, sample_rate = read_wave(wave0) + s.accept_waveform(sample_rate, samples) + + tail_paddings = np.zeros(int(0.2 * sample_rate), dtype=np.float32) + s.accept_waveform(sample_rate, tail_paddings) + + s.input_finished() + while recognizer.is_ready(s): + recognizer.decode_stream(s) + result = recognizer.get_result(s) + print(result) + self.assertTrue(len(result) > 0) + + def test_nemotron_multilingual_streaming(self): + m = "sherpa-onnx-nemotron-3.5-asr-streaming-0.6b-560ms-int8-2026-06-11" + encoder = f"{d}/{m}/encoder.int8.onnx" + decoder = f"{d}/{m}/decoder.int8.onnx" + joiner = f"{d}/{m}/joiner.int8.onnx" + tokens = f"{d}/{m}/tokens.txt" + ja_wave = f"{d}/{m}/test_wavs/ja.wav" + en_wave = f"{d}/{m}/test_wavs/en.wav" + + if not Path(encoder).is_file(): + print("skipping test_nemotron_multilingual_streaming()") + return + + recognizer = sherpa_onnx.OnlineRecognizer.from_transducer( + encoder=encoder, + decoder=decoder, + joiner=joiner, + tokens=tokens, + num_threads=1, + provider="cpu", + ) + + def decode(wave_filename: str, language: str = ""): + s = recognizer.create_stream() + if language: + s.set_option("language", language) + + samples, sample_rate = read_wave(wave_filename) + s.accept_waveform(sample_rate, samples) + + tail_paddings = np.zeros(int(0.2 * sample_rate), dtype=np.float32) + s.accept_waveform(sample_rate, tail_paddings) + + s.input_finished() + while recognizer.is_ready(s): + recognizer.decode_stream(s) + + return recognizer.get_result(s) + + ja = decode(ja_wave, "ja") + auto = decode(ja_wave, "auto") + unset = decode(ja_wave) + en = decode(en_wave, "en") + + print(f"ja: {ja}") + print(f"auto: {auto}") + print(f"unset: {unset}") + print(f"en: {en}") + + self.assertTrue(len(ja) > 0) + self.assertTrue(len(auto) > 0) + self.assertTrue(len(unset) > 0) + self.assertTrue(len(en) > 0) + self.assertTrue(any(ord(c) > 127 for c in ja)) + for result in [ja, auto, unset, en]: + self.assertIsNone(re.search(r"<[a-z]{2,3}(-[A-Z]{2})?>", result)) + def test_zipformer2_ctc(self): m = "sherpa-onnx-streaming-zipformer-ctc-multi-zh-hans-2023-12-13" for use_int8 in [True, False]: diff --git a/swift-api-examples/SherpaOnnx.swift b/swift-api-examples/SherpaOnnx.swift index a2701e157b..37c1ca3232 100644 --- a/swift-api-examples/SherpaOnnx.swift +++ b/swift-api-examples/SherpaOnnx.swift @@ -1859,6 +1859,10 @@ class SherpaOnnxOnlineStreamWrapper { } } + func setOption(key: String, value: String) { + SherpaOnnxOnlineStreamSetOption(impl, toCPointer(key), toCPointer(value)) + } + func acceptWaveform(samples: [Float], sampleRate: Int = 16000) { SherpaOnnxOnlineStreamAcceptWaveform(impl, Int32(sampleRate), samples, Int32(samples.count)) }