# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # This file was automatically generated from src/transformers/models/nemotron3_5_asr/modular_nemotron3_5_asr.py. # Do NOT edit this file manually as any edits will be overwritten by the generation of # the file from the modular. If any change should be done, please apply the change to the # modular_nemotron3_5_asr.py file directly. One of our CI enforces this. # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # Copyright 2026 The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from tokenizers.decoders import DecodeStream from ...audio_utils import AudioInput, make_list_of_audio from ...processing_utils import ProcessingKwargs, ProcessorMixin, Unpack from ...tokenization_utils_base import PreTokenizedInput, TextInput from ...utils import auto_docstring, is_torch_available, logging if is_torch_available(): import torch logger = logging.get_logger(__name__) class Nemotron3_5AsrProcessorKwargs(ProcessingKwargs, total=False): _defaults = { "audio_kwargs": { "sampling_rate": 16000, "padding": "longest", "return_attention_mask": True, "subsampling_factor": 8, }, "text_kwargs": { "padding": True, "padding_side": "right", "add_special_tokens": False, }, "common_kwargs": {"return_tensors": "pt"}, } # Default supported right attention contexts (lookaheads, in subsampled encoder frames) of the NeMo # cache-aware streaming FastConformer checkpoint. The first entry is the default. DEFAULT_NUM_LOOKAHEAD_TOKENS = [13, 6, 1, 0] DEFAULT_PROMPT_DICTIONARY = { "en-US": 0, "en": 0, "en-GB": 1, "enGB": 1, "es-ES": 2, "esES": 2, "es-US": 3, "es": 3, "zh-CN": 4, "zh-ZH": 4, "zh-TW": 5, "hi-IN": 6, "hi": 6, "hi-HI": 6, "ar-AR": 7, "ar": 7, "fr-FR": 8, "fr": 8, "de-DE": 9, "de": 9, "ja-JP": 10, "ja-JA": 10, "ru-RU": 11, "ru": 11, "pt-BR": 12, "pt-PT": 13, "pt": 13, "ko-KR": 14, "ko": 14, "ko-KO": 14, "it-IT": 15, "it": 15, "nl-NL": 16, "nl": 16, "pl-PL": 17, "pl": 17, "tr-TR": 18, "tr": 18, "uk-UA": 19, "uk": 19, "ro-RO": 20, "ro": 20, "el-GR": 21, "el": 21, "cs-CZ": 22, "cs": 22, "hu-HU": 23, "hu": 23, "sv-SE": 24, "sv": 24, "da-DK": 25, "da": 25, "fi-FI": 26, "fi": 26, "no-NO": 27, "no": 27, "nb-NO": 103, "nb": 103, "nn-NO": 104, "nn": 104, "sk-SK": 28, "sk": 28, "hr-HR": 29, "hr": 29, "bg-BG": 30, "bg": 30, "lt-LT": 31, "lt": 31, "et-EE": 60, "et": 60, "lv-LV": 61, "lv": 61, "sl-SI": 62, "sl": 62, "th-TH": 32, "vi-VN": 33, "id-ID": 34, "ms-MY": 35, "bn-IN": 36, "ur-PK": 37, "fa-IR": 38, "ta-IN": 39, "te-IN": 40, "mr-IN": 41, "gu-IN": 42, "kn-IN": 43, "ml-IN": 44, "si-LK": 45, "ne-NP": 46, "km-KH": 47, "sw-KE": 48, "am-ET": 49, "ha-NG": 50, "zu-ZA": 51, "yo-NG": 52, "ig-NG": 53, "af-ZA": 54, "rw-RW": 55, "so-SO": 56, "ny-MW": 57, "ln-CD": 58, "or-KE": 59, "he-IL": 64, "ku-TR": 65, "az-AZ": 66, "ka-GE": 67, "hy-AM": 68, "uz-UZ": 69, "tg-TJ": 70, "ky-KG": 71, "qu-PE": 80, "ay-BO": 81, "gn-PY": 82, "nah-MX": 83, "mi-NZ": 96, "haw-US": 97, "sm-WS": 98, "to-TO": 99, "fr-CA": 100, "mt-MT": 102, "auto": 101, } @auto_docstring class Nemotron3_5AsrProcessor(ProcessorMixin): def __init__( self, feature_extractor, tokenizer, blank_token="", decoder_type=None, supported_num_lookahead_tokens=None, default_num_lookahead_tokens=None, prompt_dictionary=None, num_prompts=128, ): r""" blank_token (`str`, *optional*, defaults to `""`): Blank token for RNN-T decoding. decoder_type (`str`, *optional*): Decoding/timestamp emission mode (e.g. `"ctc"`, `"rnnt"`, `"tdt"`). If `None` the decoder type is inferred automatically for backward compatibility. supported_num_lookahead_tokens (`list[int]`, *optional*): Supported right attention contexts (lookaheads, in subsampled encoder frames), mirroring `NemotronAsrStreamingEncoderConfig.supported_num_lookahead_tokens`. Used to validate `streaming_latency_ms` and to derive the returned `num_lookahead_tokens`. default_num_lookahead_tokens (`int`, *optional*): The right context used when `streaming_latency_ms` is not provided. Defaults to the first entry of `supported_num_lookahead_tokens`. prompt_dictionary (`dict[str, int]`, *optional*): Mapping from a target-language string (e.g. `"en-US"`, `"de-DE"`, the bare code `"de"`, or `"auto"`) to its prompt index. Defaults to the NeMo checkpoint's prompt dictionary. num_prompts (`int`, *optional*, defaults to 128): Number of language-prompt slots (size of the one-hot prompt vector), mirroring `Nemotron3_5AsrConfig.num_prompts`. """ self.prompt_dictionary = prompt_dictionary if prompt_dictionary is not None else DEFAULT_PROMPT_DICTIONARY self.num_prompts = num_prompts self.supported_num_lookahead_tokens = ( supported_num_lookahead_tokens if supported_num_lookahead_tokens is not None else DEFAULT_NUM_LOOKAHEAD_TOKENS ) self.default_num_lookahead_tokens = ( default_num_lookahead_tokens if default_num_lookahead_tokens is not None else self.supported_num_lookahead_tokens[0] ) self.blank_token = blank_token self.blank_token_id = tokenizer.convert_tokens_to_ids(blank_token) super().__init__(feature_extractor, tokenizer) @auto_docstring def __call__( self, audio: AudioInput, text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] | None = None, sampling_rate: int | None = None, is_streaming: bool = False, is_first_audio_chunk: bool | None = True, language: str | list[str] = "auto", **kwargs: Unpack[Nemotron3_5AsrProcessorKwargs], ): r""" sampling_rate (`int`, *optional*): The sampling rate of the input audio in Hz. Validated against the feature extractor's expected sampling rate (defaults to 16000 Hz) when provided. is_streaming (`bool`, *optional*, defaults to `False`): Whether to process audio in streaming mode (chunked), using `is_first_audio_chunk` to distinguish the first chunk from subsequent ones. is_first_audio_chunk (`bool`, *optional*, defaults to `True`): Whether the current audio is the first chunk of a streaming session. Controls `center` in the feature extractor so per-chunk STFT reproduces a single full-utterance pass. Must be `True` when `is_streaming=False`. language (`str` or `list[str]`, *optional*, defaults to `"auto"`): Target language(s) for prompt conditioning. Either a single language string applied to the whole batch, or one string per audio. Accepts locales (`"en-US"`, `"de-DE"`, ...), bare codes (`"de"`), or `"auto"` for automatic language detection. Resolved via `prompt_dictionary` into the `prompt_ids` model input. Returns: [`BatchFeature`]: the [`NemotronAsrStreamingProcessor`] outputs, augmented with: - **prompt_ids** -- A `(batch_size,)` `torch.LongTensor` of language-prompt indices. Pass it to the model/`generate`; the model turns it into the broadcast one-hot used by `prompt_projector`. """ if not is_streaming and not is_first_audio_chunk: raise ValueError("In non-streaming mode (`is_streaming=False`), `is_first_audio_chunk` must be `True`.") audio = make_list_of_audio(audio) output_kwargs = self._merge_kwargs( Nemotron3_5AsrProcessorKwargs, tokenizer_init_kwargs=self.tokenizer.init_kwargs, **kwargs, ) if sampling_rate is None: logger.warning_once( f"You've provided audio without specifying the sampling rate. It will be assumed to be " f"{output_kwargs['audio_kwargs']['sampling_rate']}, which can result in silent errors." ) elif sampling_rate != output_kwargs["audio_kwargs"]["sampling_rate"]: raise ValueError( f"The sampling rate of the audio ({sampling_rate}) does not match the sampling rate of the " f"processor ({output_kwargs['audio_kwargs']['sampling_rate']}). Please resample the audio to " f"the expected sampling rate." ) if audio is not None: # `center=True` for the first/offline chunk, `center=False` for subsequent streaming chunks. inputs = self.feature_extractor(audio, center=bool(is_first_audio_chunk), **output_kwargs["audio_kwargs"]) if text is not None: encodings = self.tokenizer(text, **output_kwargs["text_kwargs"]) inputs["num_lookahead_tokens"] = self.default_num_lookahead_tokens inputs["prompt_ids"] = self._resolve_prompt_ids(language, len(audio)) if text is None: return inputs inputs["labels"] = encodings["input_ids"] # Prepend the blank token to labels to form decoder_input_ids: the RNN-T decoder expects # [blank, label_0, ..., label_{U-1}] as input. if isinstance(text, str): text = [text] decoder_text = [self.blank_token + t for t in text] decoder_encodings = self.tokenizer(decoder_text, **output_kwargs["text_kwargs"]) inputs["decoder_input_ids"] = decoder_encodings["input_ids"] return inputs @property def model_input_names(self): feature_extractor_input_names = self.feature_extractor.model_input_names return feature_extractor_input_names + ["labels", "decoder_input_ids", "prompt_ids"] def batch_decode(self, *args, **kwargs): # RNN-T keeps repeated tokens (each is a separate emission), so consecutive identical tokens are not merged. kwargs.setdefault("group_tokens", False) return self.tokenizer.batch_decode(*args, **kwargs) def decode(self, *args, durations=None, **kwargs): """ Forward arguments to [`~PreTrainedTokenizer.decode`] and post-process the token-level timestamps (if `durations` are provided) as in the NeMo library. """ # RNN-T keeps repeated tokens (each is a separate emission), so consecutive identical tokens are not merged. kwargs.setdefault("group_tokens", False) decoded = self.tokenizer.decode(*args, **kwargs) if durations is not None: token_ids = args[0] # Derive per-step frame indices from cumulative sum of durations. timestamps = durations.cumsum(dim=-1) - durations output_kwargs = self._merge_kwargs( Nemotron3_5AsrProcessorKwargs, tokenizer_init_kwargs=self.tokenizer.init_kwargs, ) frame_rate = ( self.feature_extractor.hop_length / self.feature_extractor.sampling_rate * output_kwargs["audio_kwargs"]["subsampling_factor"] ) # Filter padding/blank tokens and decode per sequence to keep track of token-level timestamps # See `compute_rnnt_timestamps` in NeMo: # https://github.com/NVIDIA-NeMo/NeMo/blob/1692a8fb97e1aadc883cfadd2a57c4e8a1b793aa/nemo/collections/asr/parts/submodules/rnnt_decoding.py#L993 skip_ids = {self.tokenizer.pad_token_id, self.blank_token_id} proc_timestamps = [] for batch_ids, batch_timestamps in zip(token_ids, timestamps): stream = DecodeStream(skip_special_tokens=True) timestamp_dict = [] for i, token_id in enumerate(batch_ids): if int(token_id) in skip_ids: continue chunk = stream.step(self.tokenizer._tokenizer, int(token_id)) if chunk is not None: # RNN-T tokens each span a single frame (their per-step value is a 0/1 encoder advance, # not a span). start = int(batch_timestamps[i]) timestamp_dict.append( { "token": chunk, "start": start, "end": start + 1, } ) proc_timestamps.append(self._refine_timestamps(timestamp_dict, frame_rate)) return decoded, proc_timestamps return decoded def _refine_timestamps(self, char_offsets, frame_rate): # RNN-T mirrors NeMo's raw char-level timestamps, which keep every token (punctuation included) at its # own emitted frame. Only convert frame indices to seconds. for offset in char_offsets: offset["start"] = offset["start"] * frame_rate offset["end"] = offset["end"] * frame_rate return char_offsets def set_num_lookahead_tokens(self, num_lookahead_tokens: int): """ Select the right attention context (lookahead, in subsampled encoder frames) used for streaming. Sets `default_num_lookahead_tokens`, so every derived streaming property (`num_mel_frames_first_audio_chunk`, `num_mel_frames_per_audio_chunk`, `num_samples_first_audio_chunk`, `num_samples_per_audio_chunk`) re-derives from the new value. `num_lookahead_tokens` must be one of `supported_num_lookahead_tokens`. Pass the same `num_lookahead_tokens` to `model.generate` so the attention right context used in the forward matches the chunk sizes produced here; otherwise streaming `generate` raises. """ if num_lookahead_tokens not in self.supported_num_lookahead_tokens: raise ValueError( f"`num_lookahead_tokens={num_lookahead_tokens}` is not supported by this model. Supported " f"values: {list(self.supported_num_lookahead_tokens)}." ) self.default_num_lookahead_tokens = num_lookahead_tokens @property def _subsampling_factor(self) -> int: output_kwargs = self._merge_kwargs( Nemotron3_5AsrProcessorKwargs, tokenizer_init_kwargs=self.tokenizer.init_kwargs ) return output_kwargs["audio_kwargs"]["subsampling_factor"] @property def _encoder_frame_ms(self) -> float: """Duration in milliseconds of one subsampled encoder frame (`subsampling_factor * hop_length / sampling_rate`).""" return ( self._subsampling_factor * self.feature_extractor.hop_length / self.feature_extractor.sampling_rate * 1000 ) @property def streaming_latency_ms(self) -> int: """ Streaming latency (ms) of the currently-selected right attention context (`default_num_lookahead_tokens`, settable via [`~Nemotron3_5AsrProcessor.set_num_lookahead_tokens`]). The model emits a chunk only once its last frame has its full lookahead, so the delay of a right context `r` is `(r + 1)` encoder frames, i.e. `(r + 1) * encoder_frame_ms`. """ return round((self.default_num_lookahead_tokens + 1) * self._encoder_frame_ms) @property def supported_streaming_latencies_ms(self) -> dict[int, int]: """ Mapping from each supported right attention context (`supported_num_lookahead_tokens`) to its streaming latency in milliseconds (`(num_lookahead_tokens + 1) * encoder_frame_ms`). """ frame_ms = self._encoder_frame_ms return {right: round((right + 1) * frame_ms) for right in self.supported_num_lookahead_tokens} @property def num_mel_frames_first_audio_chunk(self) -> int: """ Number of mel frames the first cache-aware streaming chunk must carry, for the model's `default_num_lookahead_tokens`: `1 + subsampling_factor * num_lookahead_tokens`. """ return 1 + self._subsampling_factor * self.default_num_lookahead_tokens @property def num_mel_frames_per_audio_chunk(self) -> int: """ Number of mel frames each subsequent cache-aware streaming chunk must carry, for the model's `default_num_lookahead_tokens`: `subsampling_factor * (num_lookahead_tokens + 1)`. """ return self._subsampling_factor * (self.default_num_lookahead_tokens + 1) @property def num_samples_first_audio_chunk(self) -> int: """ Number of raw audio samples to feed the processor (with `is_first_audio_chunk=True`, i.e. `center=True`) so it returns exactly `num_mel_frames_first_audio_chunk` frames. """ return ( self.num_mel_frames_first_audio_chunk - 1 ) * self.feature_extractor.hop_length + self.feature_extractor.win_length // 2 @property def num_samples_per_audio_chunk(self) -> int: """ Number of raw audio samples to feed the processor (with `is_first_audio_chunk=False`, i.e. `center=False`) so it returns exactly `num_mel_frames_per_audio_chunk` frames. """ return ( self.num_mel_frames_per_audio_chunk * self.feature_extractor.hop_length + self.feature_extractor.win_length ) def _resolve_prompt_ids(self, language: "str | list[str]", batch_size: int) -> "torch.LongTensor": if isinstance(language, str): language = [language] * batch_size if len(language) != batch_size: raise ValueError(f"Received {len(language)} `language` entries for {batch_size} audio input(s).") prompt_ids = [] for lang in language: if lang not in self.prompt_dictionary: raise ValueError(f"Unknown `language={lang!r}`. Supported values: {sorted(self.prompt_dictionary)}.") prompt_ids.append(self.prompt_dictionary[lang]) return torch.tensor(prompt_ids, dtype=torch.long) __all__ = ["Nemotron3_5AsrProcessor"]