Quickly search /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/ for inference scripts or code that shows how to load and run the parakeet TDT 1.1B model with language tag/embedding support. Look for files like inference.py, transcribe.py, predict.py, or any code that shows model.transcribe() or language conditioning. Also check configs for model architecture details. I need to understand how to pass language tags during inference.


I'll help you search for inference scripts and code related to the Parakeet TDT 1.1B model with language tag support. Let me start by exploring the directory structure and then searching for relevant files.


Bash

find /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1 -type f -name "*.py" | grep -E "(inference|transcribe|predict)" | head -20


(Bash completed with no output)


Let me explore the directory structure first:


total 148
drwxrwxr-x 11 ubuntu ubuntu  4096 Mar 27 05:26 .
drwxrwxr-x 19 ubuntu ubuntu  4096 Mar 31 16:48 ..
-rw-rw-r--  1 ubuntu ubuntu  4939 Mar 27 05:10 analyze_tokenizer.py
drwxrwxr-x  4 ubuntu ubuntu  4096 Mar 25 02:16 configs
drwxrwxr-x  2 ubuntu ubuntu  4096 Mar 25 02:16 docs
-rw-rw-r--  1 ubuntu ubuntu   863 Mar 25 02:16 .env
-rw-rw-r--  1 ubuntu ubuntu 44503 Mar 25 02:16 finalsftdata.json
drwxrwxr-x  8 ubuntu ubuntu  4096 Mar 27 05:25 .git
-rw-rw-r--  1 ubuntu ubuntu   482 Mar 25 02:16 .gitignore
drwxrwxr-x  3 ubuntu ubuntu  4096 Mar 25 02:16 lf_asr
-rw-rw-r--  1 ubuntu ubuntu  1305 Mar 25 02:16 Makefile
-rw-rw-r--  1 ubuntu ubuntu 22260 Mar 25 02:16 PLAN.md
-rw-rw-r--  1 ubuntu ubuntu  1043 Mar 25 02:16 pyproject.toml
-rw-rw-r--  1 ubuntu ubuntu 10872 Mar 25 21:14 README.md
drwxrwxr-x  2 ubuntu ubuntu  4096 Mar 27 05:10 scripts
drwxrwxr-x  3 ubuntu ubuntu  4096 Mar 25 02:16 src
drwxrwxr-x  4 ubuntu ubuntu  4096 Mar 25 02:16 tests
drwxrwxr-x  3 ubuntu ubuntu  4096 Mar 27 05:26 tokenizers
drwxrwxr-x  2 ubuntu ubuntu  4096 Mar 25 02:16 tools


Now let me search for Python files related to inference and also explore the src directory structure:


Let me search for files containing "inference", "transcribe", or "language" in their content:


Found 4 files
maya-asr-hybrid-fastconformer-rnnt-stage1/configs/train/stage1_prod_8xh200.yaml
maya-asr-hybrid-fastconformer-rnnt-stage1/scripts/train_prod.py
maya-asr-hybrid-fastconformer-rnnt-stage1/README.md
maya-asr-hybrid-fastconformer-rnnt-stage1/PLAN.md


Let me also check the configs directory and look at the README:


find /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/configs -type f -name "*.yaml" | head -20


Now let me read the README and the main training config to understand the model architecture:


1	# Maya ASR
2
3	Multilingual Indic automatic speech recognition system built on NVIDIA NeMo. Uses a FastConformer Hybrid TDT+CTC architecture with a Parakeet TDT 1.1B pretrained encoder, fine-tuned on ~150K hours of speech across 12 languages (11 Indic + English).
4
5	## Architecture
6
7	Maya ASR is based on NeMo's `EncDecHybridRNNTCTCBPEModel` — a FastConformer encoder with dual decoding heads (TDT transducer + auxiliary CTC).
8
9	```
10	Audio (16 kHz) → Mel Spectrogram (80-dim) → FastConformer Encoder (42 layers, 1024-dim)
11	                                                      │
12	                                          ┌───────────┴───────────┐
13	                                          ▼                       ▼
14	                                   TDT Decoder+Joint         CTC Head
15	                                   (primary loss)          (auxiliary loss)
16	```
17
18	### Model Components
19
20	| Component | Architecture | Parameters |
21	|---|---|---|
22	| **Encoder** | 42-layer FastConformer (Parakeet TDT 1.1B pretrained) | ~1,062M |
23	| **TDT Decoder** | 2-layer RNN (pred_hidden=640) | ~7M |
24	| **Joint Network** | Joint hidden=640, 5 TDT duration heads | ~42M |
25	| **CTC Head** | Linear projection over encoder output | ~33M |
26	| **Language Embedding** | 12 languages → 1024-dim additive bias | ~12K |
27	| **Total** | | ~1,144M |
28
29	### Tokenizer
30
31	- 32K BPE (SentencePiece), trained on quality-filtered text from all 12 languages
32	- Located at `tokenizers/stage1_prod_bpe/`
33
34	### Supported Languages
35
36	Hindi (hi), Bengali (bn), Tamil (ta), Telugu (te), Marathi (mr), Gujarati (gu), Kannada (kn), Malayalam (ml), Punjabi (pa), Odia (or), Assamese (as), English (en).
37
38	## Training
39
40	Training is split into two stages:
41
42	### Stage 1: Acoustic Model with TDT (current)
43
44	Fine-tunes the Parakeet TDT 1.1B encoder on multilingual Indic data with a randomly initialized decoder, joint network, and CTC head sized for the 32K multilingual vocabulary.
45
46	**Key optimizations:**
47
48	- **TDT loss** — matches the pretrained encoder's original objective (durations `[0,1,2,3,4]`, sigma=0.02, omega=0.1). TDT allows the model to skip multiple frames per token emission, improving both speed and accuracy over standard RNNT.
49	- **Differential learning rate** — encoder at 2e-5 (preserves pretrained features), decoder/joint/CTC heads at 1e-3 (fast adaptation for new vocabulary).
50	- **Encoder freeze** — encoder is frozen for the first 5K optimizer steps so randomly initialized heads can warm up without corrupting pretrained weights.
51	- **CTC warmup** — CTC loss weight ramps from 0 to 0.3 over the first 3K steps, preventing random CTC gradients from interfering with early encoder adaptation.
52	- **Language embedding** — a learned per-language additive bias (12 languages to 1024-dim) is injected into the encoder output. Initialized near-zero (std=0.02) so it does not disrupt pretrained features.
53	- **Fused joint** — `fuse_loss_wer: true` with `fused_batch_size: 4` computes the joint tensor in sub-batches to cap peak VRAM.
54	- **Triple-cap OOM protection** — each micro-batch is constrained by three independent limits: max audio duration (120s), max sample count (per-bucket caps), and max total tokens (400). The token cap is critical because the TDT joint tensor size is `B x T_enc x T_dec x (V + 1 + durations)` and `T_dec` depends on transcript length, not audio length.
55
56	**Training schedule:**
57
58	| Phase | Steps | What happens |
59	|---|---|---|
60	| 0 – 3K | 0–3,000 | Encoder frozen, CTC weight ramps 0→0.3 |
61	| 3K – 5K | 3,000–5,000 | Encoder still frozen, CTC at full weight |
62	| 5K+ | 5,000–200,000 | Encoder unfrozen at 2e-5 LR, full TDT+CTC training |
63
64	**Optimizer:** AdamW (betas=0.9/0.98, weight_decay=1e-4) with CosineAnnealing schedule (2K warmup steps, min_lr=1e-5).
65
66	### Stage 2: Gemma LLM Decoder (planned)
67
68	Replaces the RNN decoder with a Gemma LLM for improved language modeling, particularly for low-resource Indic languages.
69
70	## Running Training
71
72	### Prerequisites
73
74	- 8x NVIDIA H200 GPUs (140 GB HBM3 each)
75	- Python 3.10+
76	- NeMo toolkit with ASR dependencies
77
78	### Setup
79
80	```bash
81	# Install package and dev dependencies
82	make setup
83
84	# For training dependencies (NeMo, transformers, deepspeed)
85	pip install -e ".[train]"
86	```
87
88	### Production Training (8x H200)
89
90	```bash
91	# 1. Build production data artifacts (manifests, splits, tokenizer, input config)
92	make prod-artifacts
93
94	# 2. Run full readiness check (disk, GPUs, config, artifacts, R2 storage)
95	make prod-readiness
96
97	# 3. Launch training (runs readiness gate, starts checkpoint watcher, then trains)
98	make prod-launch
99	```
100
101	Or run the training script directly:
102
103	```bash
104	python3 scripts/train_prod.py \
105	  --config configs/train/stage1_prod_8xh200.yaml \
106	  --train-parquet artifacts/phase3/production_train_final.parquet \
107	  --val-manifest data/manifests/stage1_prod_val_v2.jsonl \
108	  --pretrained-encoder pretrained/parakeet_tdt_1.1b_encoder.pt
109	```
110
111	### Smoke Test (single GPU)
112
113	```bash
114	python3 scripts/train_prod.py \
115	  --config configs/train/stage1_prod_8xh200.yaml \
116	  --max-steps 100 --devices 1 --log-every 10 --val-every 50 --smoke
117	```
118
119	### Key Training Arguments
120
121	| Argument | Default | Description |
122	|---|---|---|
123	| `--encoder-lr` | 2e-5 | Learning rate for pretrained encoder |
124	| `--head-lr` | 1e-3 | Learning rate for decoder/joint/CTC |
125	| `--freeze-encoder-steps` | 5000 | Steps to freeze encoder |
126	| `--ctc-warmup-steps` | 3000 | Steps to ramp CTC loss weight |
127	| `--max-batch-dur` | 120.0 | Max audio seconds per micro-batch |
128	| `--max-batch-size` | 16 | Hard cap on samples per micro-batch |
129	| `--max-tokens-in-batch` | 400 | Max total tokens per batch (joint tensor guard) |
130	| `--temperature` | 0.3 | Language rebalancing temperature (0=uniform, 1=natural) |
131	| `--grad-accum` | 4 | Gradient accumulation steps |
132	| `--no-lang-embed` | false | Disable language embedding |
133
134	## Inference
135
136	Maya ASR models are standard NeMo ASR models. After training produces a `.nemo` checkpoint, use NeMo's transcribe API:
137
138	```python
139	import nemo.collections.asr as nemo_asr
140
141	# Load the trained model
142	model = nemo_asr.models.EncDecHybridRNNTCTCBPEModel.restore_from("path/to/maya_asr.nemo")
143	model.eval()
144	model.cuda()
145
146	# Transcribe audio files
147	transcriptions = model.transcribe(["audio1.wav", "audio2.wav"])
148	print(transcriptions)
149	```
150
151	For TDT-specific decoding (using duration-aware greedy search):
152
153	```python
154	from omegaconf import OmegaConf
155
156	# The model defaults to TDT greedy decoding as configured in training
157	# To adjust decoding parameters:
158	decoding_cfg = OmegaConf.create({
159	    "strategy": "greedy",
160	    "model_type": "tdt",
161	    "durations": [0, 1, 2, 3, 4],
162	    "greedy": {"max_symbols": 10},
163	})
164	model.change_decoding_strategy(decoding_cfg)
165	transcriptions = model.transcribe(["audio.wav"])
166	```
167
168	To use the CTC head instead of TDT (faster but slightly less accurate):
169
170	```python
171	transcriptions = model.transcribe(["audio.wav"], return_hypotheses=False, verbose=False)
172	# Or explicitly switch to CTC decoding:
173	model.cur_decoder = "ctc"
174	transcriptions = model.transcribe(["audio.wav"])
175	```
176
177	## Data Pipeline
178
179	The training data flows through a multi-stage pipeline:
180
181	```
182	Raw audio tars (per-language shards on disk)
183	  │
184	  ▼
185	build_manifest.py → Raw manifest JSONL (path, duration, transcript, language)
186	  │
187	  ▼
188	split_manifest.py → Train/Val splits (deterministic, stratified)
189	  │
190	  ▼
191	build_tokenizer.py → 32K BPE tokenizer (SentencePiece, quality-filtered)
192	  │
193	  ▼
194	Phase 3 Production Parquet (tar_path, tar_offset_data, tar_nbytes, transcript, duration_s, language)
195	  │
196	  ▼
197	TarOffsetReader (os.pread, zero-extraction from existing tars)
198	  │
199	  ▼
200	ProductionBatchSampler (duration bucketing + language rebalancing + triple-cap OOM protection)
201	  │
202	  ▼
203	NeMo training loop
204	```
205
206	**Key data features:**
207
208	- **Tar-offset loading** — reads audio directly from tar files using byte offsets (`os.pread`), avoiding tar extraction. Zero disk overhead.
209	- **Duration bucketing** — groups similar-length clips into buckets `[0-4, 4-8, 8-12, 12-16, 16-20]s` to minimize padding waste.
210	- **Language temperature rebalancing** (T=0.3) — upweights low-resource languages so they are seen more frequently during training.
211	- **Compute-cost DDP sharding** — distributes batches across GPUs balanced by compute cost, not just sample count.
212	- **Token-length-aware batching** — pre-computes token counts per sample and caps total tokens per batch at 400, directly controlling the dominant VRAM consumer (joint tensor `T_dec` dimension).
213
214	## Hardware Requirements
215
216	### Production Training (Stage 1)
217
218	- **GPUs:** 8x NVIDIA H200 (140 GB HBM3 each)
219	- **Precision:** BF16 mixed
220	- **VRAM breakdown (per GPU):**
221	  - Model weights (bf16): ~4.5 GB
222	  - Optimizer states (AdamW, fp32 copies): ~13.5 GB
223	  - Activations + joint tensor: ~75 GB budget (75% safety margin)
224	- **Effective batch:** 8 GPUs x micro-batch x 4 gradient accumulation steps
225	- **Training target:** 200K optimizer steps, ~5-7 days
226	- **Checkpoints:** saved every 5K steps, uploaded to R2 object storage, validated (size match) before local deletion
227
228	### Inference
229
230	- **Minimum:** 1x GPU with 8+ GB VRAM (model is ~2.3 GB in bf16)
231	- **Recommended:** 1x GPU with 16+ GB VRAM for comfortable batch inference
232
233	## Project Structure
234
235	```
236	configs/
237	  train/              # Training YAML configs
238	  data/               # Data pipeline configs
239	scripts/
240	  train_prod.py       # Production training script (TDT + DiffLR + lang embed)
241	  build_manifest.py   # Raw manifest builder
242	  split_manifest.py   # Train/val splitter
243	  build_tokenizer.py  # BPE tokenizer builder
244	  run_stage1_prod.sh  # Production launcher (readiness + watcher + train)
245	lf_asr/
246	  data/
247	    production_sampler.py  # Triple-cap batch sampler with language rebalancing
248	    tar_offset_reader.py   # Zero-extraction tar audio reader
249	src/maya_asr/         # Core library
250	tokenizers/           # Trained BPE tokenizer artifacts
251	pretrained/           # Pretrained encoder weights
252	artifacts/            # Data processing outputs (parquet)
253	experiments/          # Training outputs (gitignored)
254	tests/
255	  unit/               # Unit tests (CI-safe, no data dependency)
256	  integration/        # Integration tests (requires data)
257	```
258
259	## Make Targets
260
261	| Target | Description |
262	|---|---|
263	| `make setup` | Install package + dev deps |
264	| `make test` | Run unit tests only |
265	| `make test-all` | Run unit + integration tests |
266	| `make lint` | Ruff lint check |
267	| `make format` | Auto-format + fix |
268	| `make smoke` | Import check + unit suite |
269	| `make prod-artifacts` | Build production data artifacts |
270	| `make prod-readiness` | Production launch gate (disk, GPUs, config, artifacts, R2) |
271	| `make prod-launch` | Fail-fast production launcher (readiness + watcher + training) |
272	| `make validate-r2` | End-to-end R2 upload/download/integrity test |
273
274	## License
275
276	MIT
277


Perfect! The README has inference information. Now let me check the training config and train_prod.py to see how language embeddings are used:


Error: File content (13522 tokens) exceeds maximum allowed tokens (10000). Use offset and limit parameters to read specific portions of the file, or search for specific content instead of reading the whole file.


Let me search for the language embedding implementation in train_prod.py:


54-
55-
56-class LanguageEmbedding(torch.nn.Module):
57-    """Projects a language ID into a per-frame additive bias for the encoder."""
58-
59:    def __init__(self, num_languages: int, embed_dim: int):
60-        super().__init__()
61:        self.embed = torch.nn.Embedding(num_languages, embed_dim)
62-        # Initialize near-zero so it doesn't disrupt pretrained encoder features
63-        torch.nn.init.normal_(self.embed.weight, mean=0.0, std=0.02)
64-
65:    def forward(self, lang_ids: torch.Tensor) -> torch.Tensor:
66-        """
67-        Args:
68:            lang_ids: (B,) integer tensor of language IDs
69-        Returns:
70-            (B, embed_dim, 1) — broadcast-addable to encoder output (B, D, T)
71-        """
72-        # Encoder output is (B, D, T), so we need (B, D, 1) for broadcasting
73:        return self.embed(lang_ids).unsqueeze(2)
74-
75-
76-# ---------------------------------------------------------------------------
77-# Callbacks
78-# ---------------------------------------------------------------------------
--
447-        self.transcripts = self.df["transcript"].values
448-        self.durations = self.df["duration_s"].values.astype(np.float32)
449-        self.languages = self.df["language"].values
450-
451-        # Pre-compute language IDs
452:        self.lang_ids = np.array(
453-            [LANG_TO_ID.get(str(lang).strip().lower(), LANG_TO_ID["en"]) for lang in self.languages],
454-            dtype=np.int64,
455-        )
456-
457-    def __len__(self):
--
474-        transcript = str(self.transcripts[idx]).strip()
475-        if not transcript:
476-            transcript = "<blank>"
477-        tokens = self.tokenizer.text_to_ids(transcript)
478-
479:        lang_id = int(self.lang_ids[idx])
480-
481-        return (
482-            torch.from_numpy(waveform).float(),
483-            torch.tensor(len(waveform), dtype=torch.long),
484-            torch.tensor(tokens, dtype=torch.long),
485-            torch.tensor(len(tokens), dtype=torch.long),
486:            torch.tensor(lang_id, dtype=torch.long),
487-        )
488-
489-
490-def collate_nemo_batch(batch):
491-    """Collate with padding — handles None samples from bad audio.
492-
493:    Returns (signals, sig_lens, targets, tgt_lens, lang_ids).
494-    """
495-    batch = [b for b in batch if b is not None]
496-    if not batch:
497-        return (
498-            torch.zeros(1, 1),
--
500-            torch.zeros(1, 1, dtype=torch.long),
501-            torch.tensor([0], dtype=torch.long),
502-            torch.tensor([0], dtype=torch.long),
503-        )
504-
505:    signals, sig_lens, targets, tgt_lens, lang_ids = zip(*batch)
506-
507-    max_sig_len = max(s.shape[0] for s in signals)
508-    max_tgt_len = max(t.shape[0] for t in targets)
509-
510-    padded_signals = torch.zeros(len(signals), max_sig_len)
511-    padded_targets = torch.zeros(len(targets), max_tgt_len, dtype=torch.long)
512-    signal_lengths = torch.stack(list(sig_lens))
513-    target_lengths = torch.stack(list(tgt_lens))
514:    lang_id_tensor = torch.stack(list(lang_ids))
515-
516-    for i, (sig, tgt) in enumerate(zip(signals, targets)):
517-        padded_signals[i, :sig.shape[0]] = sig
518-        padded_targets[i, :tgt.shape[0]] = tgt
519-
520:    return padded_signals, signal_lengths, padded_targets, target_lengths, lang_id_tensor
521-
522-
523-def get_ddp_info() -> tuple[int, int]:
524-    """Get (rank, world_size) from DDP environment."""
525-    if torch.distributed.is_available() and torch.distributed.is_initialized():
--
530-
531-
532-# ---------------------------------------------------------------------------
533-# Language-conditioned forward hook
534-# ---------------------------------------------------------------------------
535:def install_language_conditioning(model, lang_embed: LanguageEmbedding):
536:    """Install a forward hook on the encoder to add language embeddings.
537-
538:    After the encoder produces its output, we add the language embedding
539-    as a per-frame bias. This gives the encoder a language signal without
540-    modifying the encoder architecture (preserving pretrained weights).
541-    """
542-    # Store reference on model so it moves to GPU with model
543:    model._lang_embed = lang_embed
544-
545-    original_forward = model.encoder.forward
546-
547:    def encoder_forward_with_lang(audio_signal, length, lang_ids=None):
548-        encoded, encoded_len = original_forward(audio_signal=audio_signal, length=length)
549:        if lang_ids is not None and hasattr(model, '_lang_embed'):
550:            lang_bias = model._lang_embed(lang_ids)  # (B, 1, D)
551-            encoded = encoded + lang_bias
552-        return encoded, encoded_len
553-
554-    model.encoder.forward = encoder_forward_with_lang
555-    return model
--
557-
558-# ---------------------------------------------------------------------------
559-# Custom training step with language conditioning
560-# ---------------------------------------------------------------------------
561-def patch_training_step(model):
562:    """Patch the model's training_step to extract lang_ids from batch
563-    and pass them through to the encoder via language conditioning.
564-    """
565-    original_training_step = model.training_step
566-
567-    def patched_training_step(batch, batch_idx):
568-        # NeMo expects batch = (signal, signal_len, tokens, token_len)
569:        # We added lang_ids as 5th element
570-        if isinstance(batch, (list, tuple)) and len(batch) >= 5:
571:            signal, signal_len, tokens, token_len, lang_ids = batch[:5]
572:            # Store lang_ids on model for encoder hook to pick up
573:            model._current_lang_ids = lang_ids.to(signal.device)
574-            # Pass standard 4-tuple to NeMo
575-            nemo_batch = (signal, signal_len, tokens, token_len)
576-            return original_training_step(nemo_batch, batch_idx)
577-        return original_training_step(batch, batch_idx)
578-
--
583-    """Patch encoder forward to inject language IDs stored on model."""
584-    original_encoder_forward = model.encoder.forward
585-
586-    def encoder_forward_with_lang(audio_signal, length):
587-        encoded, encoded_len = original_encoder_forward(audio_signal=audio_signal, length=length)
588:        if (hasattr(model, '_current_lang_ids') and hasattr(model, '_lang_embed')
589:                and model._current_lang_ids is not None):
590:            lang_ids = model._current_lang_ids
591-            # Safety: only apply if batch sizes match (skip during validation)
592:            if lang_ids.shape[0] == encoded.shape[0]:
593:                lang_bias = model._lang_embed(lang_ids)  # (B, D, 1)
594-                encoded = encoded + lang_bias
595-            # Clear after use to prevent stale values during validation
596:            model._current_lang_ids = None
597-        return encoded, encoded_len
598-
599-    model.encoder.forward = encoder_forward_with_lang
600-
601-
--
649-    parser.add_argument("--freeze-encoder-steps", type=int, default=5000,
650-                        help="Freeze encoder for this many optimizer steps (0=no freeze)")
651-    parser.add_argument("--ctc-warmup-steps", type=int, default=3000,
652-                        help="Steps to warm up CTC loss from 0 to target weight")
653-    parser.add_argument("--no-lang-embed", action="store_true", default=False,
654:                        help="Disable language embedding conditioning")
655-    args = parser.parse_args()
656-
657-    cfg = OmegaConf.load(args.config)
658-
659-    with open_dict(cfg):
--
719-        print(f"  Loss:           TDT (durations=[0,1,2,3,4])")
720-        print(f"  Encoder LR:     {args.encoder_lr}")
721-        print(f"  Head LR:        {args.head_lr}")
722-        print(f"  Encoder freeze: {args.freeze_encoder_steps} steps")
723-        print(f"  CTC warmup:     0→target over {args.ctc_warmup_steps} steps")
724:        print(f"  Lang embed:     {'yes' if not args.no_lang_embed else 'no'}")
725-        print()
726-        print("VRAM Budget (OOM-safe):")
727-        print(f"  GPU memory:     {vram_budget['gpu_memory_gb']:.0f} GB")
728-        print(f"  Model (bf16):   {vram_budget['model_params_gb']:.1f} GB")
729-        print(f"  Optimizer:      {vram_budget['optimizer_gb']:.1f} GB")
--
797-        print()
798-
799-    # -----------------------------------------------------------------------
800-    # Language conditioning
801-    # -----------------------------------------------------------------------
802:    if not args.no_lang_embed:
803:        lang_embed = LanguageEmbedding(NUM_LANGUAGES, cfg.model.model_defaults.enc_hidden)
804:        model._lang_embed = lang_embed
805-        # Register as a proper submodule so it moves to GPU and is saved
806:        model.register_module('_lang_embed_module', lang_embed)
807-        patch_encoder_for_lang(model)
808-        patch_training_step(model)
809-        if rank == 0:
810:            lang_params = sum(p.numel() for p in lang_embed.parameters())
811-            print(f"  [LangEmbed] {NUM_LANGUAGES} languages → {cfg.model.model_defaults.enc_hidden}D ({lang_params:,} params)")
812-
813-    # -----------------------------------------------------------------------
814-    # Callbacks
815-    # -----------------------------------------------------------------------


Great! Now let me check the LANG_TO_ID mapping and look at the full context of how inference would use language tags:


46-# ---------------------------------------------------------------------------
47-# Language embedding module
48-# ---------------------------------------------------------------------------
49:LANG_TO_ID = {
50-    "hi": 0, "bn": 1, "ta": 2, "te": 3, "mr": 4, "gu": 5,
51-    "kn": 6, "ml": 7, "pa": 8, "or": 9, "as": 10, "en": 11,
52-}
53:NUM_LANGUAGES = len(LANG_TO_ID)
54-
55-
56-class LanguageEmbedding(torch.nn.Module):
--
450-
451-        # Pre-compute language IDs
452-        self.lang_ids = np.array(
453:            [LANG_TO_ID.get(str(lang).strip().lower(), LANG_TO_ID["en"]) for lang in self.languages],
454-            dtype=np.int64,
455-        )
456-
--
800-    # Language conditioning
801-    # -----------------------------------------------------------------------
802-    if not args.no_lang_embed:
803:        lang_embed = LanguageEmbedding(NUM_LANGUAGES, cfg.model.model_defaults.enc_hidden)
804-        model._lang_embed = lang_embed
805-        # Register as a proper submodule so it moves to GPU and is saved
806-        model.register_module('_lang_embed_module', lang_embed)
--
808-        patch_training_step(model)
809-        if rank == 0:
810-            lang_params = sum(p.numel() for p in lang_embed.parameters())
811:            print(f"  [LangEmbed] {NUM_LANGUAGES} languages → {cfg.model.model_defaults.enc_hidden}D ({lang_params:,} params)")
812-
813-    # -----------------------------------------------------------------------
814-    # Callbacks


Perfect! Now let me read a bit more of the train_prod.py to see the encoder patching logic and then check the config file:


1	#!/usr/bin/env python3
2	"""Production training script with optimized data loading + TDT + differential LR.
3
4	Data strategy (ported from parakeet pipeline):
5	  - Duration bucketing: groups similar-length clips → minimal padding waste
6	  - Language temperature rebalancing (T=0.3): upweights low-resource languages
7	  - Max-batch-duration packing: caps audio per micro-batch → stable VRAM
8	  - Compute-cost DDP sharding: balances load across GPUs
9	  - Grad accum alignment: prevents DDP sync bugs
10
11	Optimizations v2:
12	  - TDT loss (matching pretrained Parakeet encoder objective)
13	  - Differential LR: encoder=2e-5, decoder/joint/CTC=1e-3
14	  - Encoder frozen for first 5K steps, then unfrozen
15	  - CTC loss warmup: 0→0.3 over first 3K steps
16	  - Language embedding conditioning
17	  - OOM-safe VRAM ceiling design
18
19	Audio I/O: TarOffsetReader (os.pread, zero extraction from existing tars).
20	Model: NeMo EncDecHybridRNNTCTCBPEModel (FastConformer + TDT + aux CTC).
21
22	Usage:
23	  # Full production run (8×H200)
24	  python3 scripts/train_prod.py --config configs/train/stage1_prod_8xh200.yaml
25
26	  # Quick smoke test
27	  python3 scripts/train_prod.py --config configs/train/stage1_prod_8xh200.yaml \
28	    --max-steps 100 --devices 1 --log-every 10 --val-every 50
29	"""
30
31	import argparse
32	import os
33	import sys
34	import time
35	from pathlib import Path
36
37	import lightning.pytorch as pl
38	import nemo.collections.asr as nemo_asr
39	import numpy as np
40	import pyarrow.parquet as pq
41	import torch
42	from nemo.utils.exp_manager import exp_manager
43	from omegaconf import OmegaConf, open_dict
44
45
46	# ---------------------------------------------------------------------------
47	# Language embedding module
48	# ---------------------------------------------------------------------------
49	LANG_TO_ID = {
50	    "hi": 0, "bn": 1, "ta": 2, "te": 3, "mr": 4, "gu": 5,
51	    "kn": 6, "ml": 7, "pa": 8, "or": 9, "as": 10, "en": 11,
52	}
53	NUM_LANGUAGES = len(LANG_TO_ID)
54
55
56	class LanguageEmbedding(torch.nn.Module):
57	    """Projects a language ID into a per-frame additive bias for the encoder."""
58
59	    def __init__(self, num_languages: int, embed_dim: int):
60	        super().__init__()
61	        self.embed = torch.nn.Embedding(num_languages, embed_dim)
62	        # Initialize near-zero so it doesn't disrupt pretrained encoder features
63	        torch.nn.init.normal_(self.embed.weight, mean=0.0, std=0.02)
64
65	    def forward(self, lang_ids: torch.Tensor) -> torch.Tensor:
66	        """
67	        Args:
68	            lang_ids: (B,) integer tensor of language IDs
69	        Returns:
70	            (B, embed_dim, 1) — broadcast-addable to encoder output (B, D, T)
71	        """
72	        # Encoder output is (B, D, T), so we need (B, D, 1) for broadcasting
73	        return self.embed(lang_ids).unsqueeze(2)
74
75
76	# ---------------------------------------------------------------------------
77	# Callbacks
78	# ---------------------------------------------------------------------------
79	class ThroughputCallback(pl.Callback):
80	    """Logs VRAM, throughput (samples/s, audio-s/s), and MFU every log_every steps."""
81
82	    def __init__(self, log_every: int = 10, model_params: int = 0,
83	                 grad_accum: int = 1, world_size: int = 1):
84	        self.log_every = log_every
85	        self.model_params = model_params
86	        self.grad_accum = grad_accum
87	        self.world_size = world_size
88	        self._window_t0 = None
89	        self._window_samples = 0
90	        self._window_audio_s = 0.0
91	        self._window_micro_steps = 0
92	        # H200 specs: bf16 peak ~990 TFLOPS
93	        self._gpu_peak_tflops = 989.0
94
95	    def on_train_batch_start(self, trainer, pl_module, batch, batch_idx):
96	        if self._window_t0 is None:
97	            self._window_t0 = time.time()
98
99	    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
100	        if batch is not None and len(batch) >= 2:


Now let me check the config file to understand the model architecture better:


1	# Stage 1 Production Config: FastConformer Hybrid TDT+CTC on full dataset
2	# Hardware: 8x NVIDIA H200 (140 GB HBM3 each)
3	# Target: ~200K steps, ~5-7 days
4	#
5	# Pretrained encoder: Parakeet TDT 1.1B (42-layer FastConformer, English)
6	# Tokenizer: encoder-side 32K BPE (quality-filtered text, 11 Indic + English)
7	# Decoder/joint: randomly initialized for 32K vocab, then finetuned
8	#
9	# Optimizations v2:
10	#   - TDT loss (matching pretrained encoder objective) instead of RNNT
11	#   - Differential LR: encoder=2e-5, decoder/joint/CTC=1e-3
12	#   - Encoder frozen for first 5K steps, then unfrozen at low LR
13	#   - CTC loss warmup: 0 -> 0.3 over first 3K steps
14	#   - Language embedding conditioning (lang_id -> encoder)
15
16	name: maya-asr-stage1-prod-v2
17
18	trainer:
19	  devices: 8
20	  num_nodes: 1
21	  accelerator: gpu
22	  strategy: ddp_find_unused_parameters_true
23	  precision: bf16-mixed
24	  max_steps: 200000
25	  val_check_interval: 10000
26	  accumulate_grad_batches: 4  # effective = 8 GPUs x batch x 4 accum
27	  gradient_clip_val: 1.0
28	  log_every_n_steps: 100
29	  enable_checkpointing: false  # handled by exp_manager
30	  default_root_dir: experiments/stage1_prod
31
32	model:
33	  model_defaults:
34	    asr_enc_hidden: 1024
35	    enc_hidden: 1024
36	    dec_hidden: 640
37	    pred_hidden: 640
38
39	  preprocessor:
40	    _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor
41	    sample_rate: 16000
42	    normalize: per_feature
43	    window_size: 0.025
44	    window_stride: 0.01
45	    features: 80
46	    n_fft: 512
47	    dither: 0.00001
48
49	  encoder:
50	    _target_: nemo.collections.asr.modules.ConformerEncoder
51	    feat_in: 80
52	    n_layers: 42
53	    d_model: ${model.model_defaults.enc_hidden}
54	    subsampling: dw_striding
55	    subsampling_factor: 8
56	    subsampling_conv_channels: 256
57	    ff_expansion_factor: 4
58	    self_attention_model: rel_pos
59	    n_heads: 8
60	    conv_kernel_size: 9
61	    dropout: 0.1
62	    dropout_pre_encoder: 0.1
63	    dropout_emb: 0.0
64
65	  decoder:
66	    _target_: nemo.collections.asr.modules.RNNTDecoder
67	    prednet:
68	      pred_hidden: ${model.model_defaults.pred_hidden}
69	      pred_rnn_layers: 2
70	      dropout: 0.1
71
72	  joint:
73	    _target_: nemo.collections.asr.modules.RNNTJoint
74	    num_extra_outputs: 5  # TDT duration heads (durations=[0,1,2,3,4])
75	    fuse_loss_wer: true
76	    fused_batch_size: 2   # sub-batch joint computation → caps peak VRAM
77	    jointnet:
78	      joint_hidden: 640
79	      activation: relu
80	      dropout: 0.1
81	      encoder_hidden: ${model.model_defaults.enc_hidden}
82	      pred_hidden: ${model.model_defaults.pred_hidden}
83
84	  loss:
85	    loss_name: tdt
86	    tdt_kwargs:
87	      durations: [0, 1, 2, 3, 4]
88	      sigma: 0.02
89	      omega: 0.1
90	      fastemit_lambda: 0.0
91	      clamp: -1.0
92
93	  decoding:
94	    strategy: greedy
95	    model_type: tdt
96	    durations: [0, 1, 2, 3, 4]
97	    greedy:
98	      max_symbols: 10
99
100	  aux_ctc:
101	    ctc_loss_weight: 0.3  # warmed up from 0->0.3 over 3K steps in train_prod.py
102	    use_cer: false
103	    decoder:
104	      _target_: nemo.collections.asr.modules.ConvASRDecoder
105	      feat_in: null
106	      num_classes: -1
107
108	  # Encoder-side tokenizer (32K BPE, quality-filtered, 11 Indic + English)
109	  tokenizer:
110	    dir: tokenizers/stage1_prod_bpe
111	    type: bpe
112
113	  train_ds:
114	    # Production: tar-offset loading via train_prod.py (not NeMo tarred format)
115	    manifest_filepath: data/manifests/stage1_prod_val_v2.jsonl  # placeholder for model init
116	    batch_size: 4
117	    shuffle: true
118	    num_workers: 8
119	    max_duration: 20.0
120	    min_duration: 0.5
121	    sample_rate: 16000
122	    trim_silence: false
123
124	  validation_ds:
125	    manifest_filepath: data/manifests/stage1_prod_val_v2.jsonl
126	    batch_size: 4
127	    shuffle: false
128	    num_workers: 4
129	    sample_rate: 16000
130
131	  # NOTE: optim LR is overridden by differential LR in train_prod.py
132	  # This base LR is used for NeMo's scheduler init but actual param group
133	  # LRs are set in the training script.
134	  optim:
135	    name: adamw
136	    lr: 0.001  # decoder/joint/CTC LR (encoder uses 2e-5 via param groups)
137	    betas: [0.9, 0.98]
138	    weight_decay: 0.0001
139	    sched:
140	      name: CosineAnnealing
141	      warmup_steps: 2000
142	      min_lr: 0.00001
143
144	exp_manager:
145	  exp_dir: experiments/stage1_prod
146	  name: maya_asr_stage1_tdt_v2
147	  create_tensorboard_logger: true
148	  create_wandb_logger: true
149	  wandb_logger_kwargs:
150	    project: maya-asr
151	    name: tdt-parakeet-finetune-stage1-v2
152	    resume: allow
153	  resume_if_exists: false
154	  resume_ignore_no_checkpoint: true
155	  create_checkpoint_callback: true
156	  checkpoint_callback_params:
157	    monitor: val_loss
158	    mode: min
159	    save_top_k: 3
160	    save_last: true
161	    every_n_train_steps: 5000
162	    every_n_epochs: 0
163	    save_nemo_on_train_end: true
164


Now let me search for any example inference code or additional documentation about how to use language tags at inference time:


Found 3 files
maya-asr-hybrid-fastconformer-rnnt-stage1/scripts/train_prod.py
maya-asr-hybrid-fastconformer-rnnt-stage1/README.md
maya-asr-hybrid-fastconformer-rnnt-stage1/PLAN.md


Let me check the PLAN.md file for any additional details about inference:


1-# Maya ASR: FastConformer + Gemma Transcription Build Guide
2-
3-## Executive Summary
4-
5:Build a speech-to-text transcription system using a FastConformer encoder coupled with a small Gemma 3 decoder, trained at full-parameter scale on ~150K hours of transcribed Indic+English data across 12 languages. **Task: transcription only** (no translation). Covers architecture design, tokenizer construction, multi-stage training recipes, distributed training on 8×H100 GPUs, and concrete NeMo scripts and config references.
6-
7----
8-
9-## 1. Architecture Overview
--
21-│  │  8x subsampling       │    │ MLP      │    │              │   │
22-│  └───────────────────────┘    └──────────┘    └──────────────┘   │
23-│                                                                  │
24-│  Language Prompting: <|source_lang|> token guides decoder        │
25:│  Task: <|transcribe|> only                                       │
26-└──────────────────────────────────────────────────────────────────┘
27-```
28-
29-### 1.2 Component Breakdown
30-
31-**Encoder: FastConformer (IndicConformer-600M base)**
32-
33:Based on [AI4Bharat's IndicConformer-600M multilingual model](https://huggingface.co/ai4bharat/indic-conformer-600m-multilingual), a Conformer-based hybrid CTC+RNNT model supporting all 22 scheduled Indian languages.
34-
35-- Architecture: FastConformer with 8x depthwise convolutional subsampling (256 channels)
36-- Reduced convolutional kernel size of 9 in conformer blocks
37:- 2.4x faster training/inference than standard Conformer without quality degradation
38-- Supports local attention for audio >1 hour
39-- NeMo-native: `EncDecHybridRNNTCTCBPEModel` class
40-
41-Scale to **32 encoder layers at 1024 hidden dimension** (matching Canary-1B-Flash's encoder spec). Initialize from IndicConformer-600M checkpoint and expand layers if needed.
--
69-| Telugu | 2.93 | 404 tokens |
70-| Kannada | 3.15 | 231 tokens |
71-| Malayalam | 3.39 | ~200 tokens (est.) |
72-
73:With 150K hours, full instruction tuning (not LoRA) is the right approach — the decoder needs to deeply learn acoustic-to-text alignment for all 12 languages.
74-
75-If Gemma 3 1B proves too heavy, consider distilling down or using only a subset of its layers (e.g., first 16 of 26 layers) as a smaller decoder, keeping the tokenizer.
76-
77-**Projection Layer: Linear MLP**
--
110-
111-Organized as `lang={code}/{shard_id}/` with each shard containing:
112-- `audio.tar` — compressed FLAC audio files (~15,000 segments per shard)
113-- `audio_index.parquet` — audio indexing metadata
114:- `manifest.json` — shard-level metadata (counts, checksums, language)
115-- `metadata.parquet` — per-segment metadata (transcription, duration, speaker, domain)
116-- `xcodec2_tokens.parquet` — pre-computed acoustic tokens (optional)
117-
118:**Shard counts per language (final-export):**
119-
120-| Language | Shards | Est. Hours |
121-|---|---|---|
122-| English (en) | 840 | ~28K |
--
164-
165-```yaml
166-input_cfg:
167-  - corpus: final_export_hindi
168:    language: hi
169-    manifest_filepath: /data/nemo_manifests/hi/manifest.json
170-    tarred_audio_filepaths: /data/nemo_tarred/hi/audio_*.tar
171-    type: nemo_tarred
172-    weight: 10000
173-  - corpus: final_export_bengali
174:    language: bn
175-    manifest_filepath: /data/nemo_manifests/bn/manifest.json
176-    tarred_audio_filepaths: /data/nemo_tarred/bn/audio_*.tar
177-    type: nemo_tarred
178-    weight: 10000
179:  # ... repeat for all 12 languages + supplementary sources
180-```
181-
182-### 2.6 Language Balancing
183-
184-Use **two-stage deterministic upsampling** ([Polyglot-Lion, March 2026](https://arxiv.org/html/2603.16184v1)):
185-
186:1. **Intra-language**: Upsample smaller corpora within each language to match the largest corpus
187:2. **Inter-language**: Upsample each language to match the largest language group
188-
189:This produced 72% relative WER reduction for underrepresented languages without degrading high-resource ones.
190-
191-Alternative: Temperature-based sampling (α = 0.3–0.5) as used by [Canary-1B-v2](https://arxiv.org/html/2509.14128v2).
192-
193-### 2.7 Data Augmentation
--
206-
207-Use Gemma 3's existing 262K SentencePiece tokenizer directly. It already has strong Indic coverage with 4K+ dedicated Indic tokens. No custom tokenizer training needed for the decoder side.
208-
209-Add special tokens:
210:- `<|transcribe|>` — task token
211-- Language ID tokens: `<|hi|>`, `<|bn|>`, `<|ta|>`, `<|te|>`, `<|gu|>`, `<|kn|>`, `<|ml|>`, `<|mr|>`, `<|pa|>`, `<|or|>`, `<|as|>`, `<|en|>`
212-
213-### 3.2 Encoder-Side Tokenizer (for CTC Auxiliary Loss)
214-
--
312-for batch in dataloader:
313-    audio_features = encoder(batch.audio)       # frozen
314-    projected = projection(audio_features)       # trainable
315-
316:    # Prepend language + task tokens: <|hi|><|transcribe|>
317-    prompt_embeds = llm.embed_tokens(batch.prompt_ids)
318-    input_embeds = torch.cat([projected, prompt_embeds], dim=1)
319-
320-    logits = llm(inputs_embeds=input_embeds)     # frozen
--
335-Code reference: [github.com/csalt-research/salsa](https://github.com/csalt-research/salsa)
336-
337-### 4.4 Stage 3: Quality Fine-Tuning
338-
339:Fine-tune on a **language-balanced, high-quality subset**:
340-
341:- Select cleanest 10% of data per language
342-- Train for 10K–20K additional steps
343-- Lower learning rate (0.1× of Stage 2 peak)
344-- This stage polishes accuracy and reduces hallucination
345-
--
349-
350-### 5.1 Hardware: 8×H100 80GB
351-
352-For a ~800M–1.2B parameter model (600M encoder + projection + Gemma decoder):
353:- **Single H100**: 80GB VRAM, sufficient for inference but tight for full training
354-- **8×H100**: DDP as default for encoder-only stages. DeepSpeed ZeRO-2 for joint training with Gemma
355-
356-### 5.2 NeMo Trainer Config (Encoder Stages)
357-
--
417-    max_duration: 40.0
418-    min_duration: 0.01
419-  ```
420-  Use NeMo's [2D duration bin estimation script](https://github.com/NVIDIA-NeMo/NeMo/blob/main/scripts/speech_recognition/estimate_duration_bins_2d.py)
421:- **Tarred audio datasets**: Pre-tar into WebDataset shards (512–1024 per language). Our data is already shard-tarred.
422-- **Gradient accumulation**: accumulate_grad_batches=4–8 for large effective batch sizes
423-- **Cosine annealing scheduler**: Standard for FastConformer training
424-
425----
--
437-| [Kathbath](https://github.com/AI4Bharat/vistaar) | 12 Indic | Conversational | AI4Bharat |
438-
439-### 6.2 Metrics
440-
441:- **WER** (Word Error Rate): Primary metric for all languages
442:- **CER** (Character Error Rate): Secondary for morphologically rich languages (Tamil, Malayalam, Telugu)
443-- **RTFx** (Real-Time Factor, inverse): Throughput measurement
444:- Per-language breakdown — never average across languages without reporting individual scores
445-
446-### 6.3 Baselines to Compare Against
447-
448-| Model | Best Hindi WER | Notes |
--
455----
456-
457-## 7. Avoiding Catastrophic Forgetting
458-
459:With 12 languages trained jointly, forgetting is less of a concern. Key strategies:
460-
461:- **Joint training** (all languages simultaneously): Best approach. No forgetting by design.
462:- **Language-balanced sampling**: Critical — without it, high-resource languages dominate
463:- **If adding languages later**: Use Learning without Forgetting (LwF) — consistently low and stable WER, best BWT in noisy settings
464:- **CTC auxiliary loss**: Acts as implicit regularizer, stabilizing encoder representations across languages
465-
466-Code: [github.com/FrozenWolf-Cyber/Indic-CL-ASR](https://github.com/FrozenWolf-Cyber/Indic-CL-ASR)
467-
468----
--
474-| Script | Path | Purpose |
475-|---|---|---|
476-| Hybrid CTC+RNNT Training | `examples/asr/asr_hybrid_transducer_ctc/speech_to_text_hybrid_rnnt_ctc_bpe_prompt.py` | Encoder training with prompt support |
477-| AED Multitask Training | `examples/asr/speech_multitask/speech_to_text_aed.py` | Canary-style training (reference) |
478:| CTC Language Fine-tuning | `tutorials/asr/ASR_CTC_Language_Finetuning.ipynb` | Tutorial for adding new languages |
479-| Tokenizer Training | `scripts/tokenizers/process_asr_text_tokenizer.py` | Build SentencePiece tokenizer |
480-| Duration Bin Estimation | `scripts/speech_recognition/estimate_duration_bins_2d.py` | Optimize Lhotse bucketing |
481-| SALSA/SKIP-SALSA | `github.com/csalt-research/salsa` | ASR-LLM synchronous coupling |
482-
--
493-### 8.3 Key Papers
494-
495-- [Canary-1B-v2 & Parakeet-TDT-0.6B-v3 Technical Report](https://arxiv.org/html/2509.14128v2) — NVIDIA's training methodology
496-- [SALSA: Speedy ASR-LLM Synchronous Aggregation](https://arxiv.org/abs/2408.16542) — ASR+LLM coupling via projection layers
497:- [SKIP-SALSA](https://www.isca-archive.org/interspeech_2025/mittal25_interspeech.pdf) — Handles token fertility gap for Indic languages
498-- [Bridging the Gap: Speech-LLM vs E2E ASR](https://arxiv.org/html/2601.01461v1) — Two-stage training beats tri-stage
499-- [Polyglot-Lion](https://arxiv.org/html/2603.16184v1) — Balanced multilingual sampling
500:- [Indic CL-ASR](https://arxiv.org/abs/2508.06280) — Continual learning for Indian language ASR
501-- [Slam: Training SLMs on One GPU in a Day](https://arxiv.org/html/2502.15814v1) — Efficient speech-LLM training
502-- [Vistaar/IndicWhisper](https://arxiv.org/abs/2305.15386) — Indic training data and benchmarks
503-
504-### 8.4 NeMo Installation


Let me also look at how the model is initialized and check if there's any example of inference setup:


750	    model = nemo_asr.models.EncDecHybridRNNTCTCBPEModel(cfg=cfg.model, trainer=trainer)
751	    params = sum(p.numel() for p in model.parameters())
752	    if rank == 0:
753	        print(f"Model parameters: {params:,}")
754
755	    # Resume weights from checkpoint (without optimizer state)
756	    if args.resume_weights_only and args.resume_from_checkpoint:
757	        ckpt_path = args.resume_from_checkpoint
758	        if rank == 0:
759	            print(f"\nResuming weights only from: {ckpt_path}")
760	        ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
761	        model.load_state_dict(ckpt["state_dict"], strict=False)
762	        resume_step = ckpt.get("global_step", 0)
763	        if rank == 0:
764	            print(f"  Loaded model weights from step {resume_step}")
765	            print(f"  Optimizer state NOT restored (fresh optimizer with DiffLR)")
766	        # Clear resume flag so Lightning doesn't also try to restore
767	        args.resume_from_checkpoint = None
768	        with open_dict(cfg):
769	            cfg.exp_manager.resume_from_checkpoint = None
770	            cfg.exp_manager.resume_if_exists = False
771	        del ckpt
772	        print()
773
774	    # Load pretrained encoder weights
775	    if args.pretrained_encoder is not None and not args.resume_weights_only:
776	        if rank == 0:
777	            print(f"\nLoading pretrained encoder from: {args.pretrained_encoder}")
778	        pretrained_state = torch.load(args.pretrained_encoder, map_location="cpu", weights_only=True)
779	        model_state = model.state_dict()
780
781	        loaded, skipped = 0, 0
782	        for key, value in pretrained_state.items():
783	            if key in model_state and model_state[key].shape == value.shape:
784	                model_state[key] = value
785	                loaded += 1
786	            else:
787	                skipped += 1
788	                if rank == 0 and skipped <= 10:
789	                    expected = model_state[key].shape if key in model_state else "NOT FOUND"
790	                    print(f"  Skip: {key} pretrained={value.shape} model={expected}")
791
792	        model.load_state_dict(model_state)
793	        enc_params = sum(pretrained_state[k].numel() for k in pretrained_state if k in model_state and model_state[k].shape == pretrained_state[k].shape)
794	        if rank == 0:
795	            print(f"  Loaded {loaded} tensors ({enc_params/1e6:.0f}M params), skipped {skipped}")
796	            print(f"  Decoder + joint + CTC head: randomly initialized for {model.tokenizer.vocab_size} vocab")
797	        print()
798
799	    # -----------------------------------------------------------------------
800	    # Language conditioning
801	    # -----------------------------------------------------------------------
802	    if not args.no_lang_embed:
803	        lang_embed = LanguageEmbedding(NUM_LANGUAGES, cfg.model.model_defaults.enc_hidden)
804	        model._lang_embed = lang_embed
805	        # Register as a proper submodule so it moves to GPU and is saved
806	        model.register_module('_lang_embed_module', lang_embed)
807	        patch_encoder_for_lang(model)
808	        patch_training_step(model)
809	        if rank == 0:
810	            lang_params = sum(p.numel() for p in lang_embed.parameters())
811	            print(f"  [LangEmbed] {NUM_LANGUAGES} languages → {cfg.model.model_defaults.enc_hidden}D ({lang_params:,} params)")
812
813	    # -----------------------------------------------------------------------
814	    # Callbacks
815	    # -----------------------------------------------------------------------
816	    log_interval = args.log_every or cfg.trainer.get("log_every_n_steps", 100)
817	    throughput_cb = ThroughputCallback(
818	        log_every=log_interval,
819	        model_params=params,
820	        grad_accum=grad_accum,
821	        world_size=world_size,
822	    )
823	    trainer.callbacks.append(throughput_cb)
824
825	    freeze_cb = EncoderFreezeCallback(freeze_steps=args.freeze_encoder_steps)
826	    trainer.callbacks.append(freeze_cb)
827
828	    ctc_warmup_cb = CTCWarmupCallback(
829	        target_weight=cfg.model.aux_ctc.get("ctc_loss_weight", 0.3),
830	        warmup_steps=args.ctc_warmup_steps,
831	    )
832	    trainer.callbacks.append(ctc_warmup_cb)
833
834	    diff_lr_cb = DifferentialLRCallback(
835	        encoder_lr=args.encoder_lr,
836	        head_lr=args.head_lr,
837	    )
838	    trainer.callbacks.append(diff_lr_cb)
839
840	    # -----------------------------------------------------------------------
841	    # Load train dataset
842	    # -----------------------------------------------------------------------
843	    if rank == 0:
844	        print(f"Loading train parquet: {args.train_parquet}")
845	    train_dataset = NeMoTarOffsetDataset(
846	        manifest_parquet=args.train_parquet,
847	        tokenizer=model.tokenizer,
848	        sample_rate=cfg.model.train_ds.get("sample_rate", 16000),
849	        max_duration=max_dur,
850	        min_duration=min_dur,
851	    )
852	    if rank == 0:
853	        print(f"Train samples: {len(train_dataset):,}")
854	        print(f"Train audio:   {train_dataset.durations.sum() / 3600:.0f} hours")
855	        # Language distribution
856	        unique_langs, lang_counts = np.unique(train_dataset.languages, return_counts=True)
857	        print(f"Languages:     {len(unique_langs)} ({', '.join(sorted(unique_langs))})")
858	        print()
859
860	    # Create production batch sampler
861	    # Pre-compute token lengths for VRAM-safe batching
862	    # Token length controls T_dec in the joint tensor — the KEY factor for OOM
863	    # Cached to .npy file so subsequent runs skip tokenization entirely
864	    cache_path = Path(args.train_parquet).with_suffix(".token_lengths.npy")
865	    if cache_path.exists():
866	        if rank == 0:
867	            print(f"Loading cached token lengths from {cache_path}")
868	        token_lengths = np.load(cache_path)
869	        if len(token_lengths) != len(train_dataset):
870	            if rank == 0:
871	                print(f"  Cache size mismatch ({len(token_lengths)} vs {len(train_dataset)}), recomputing...")
872	            cache_path.unlink()
873	            token_lengths = None
874	        else:
875	            if rank == 0:
876	                print(f"  Loaded {len(token_lengths):,} cached token lengths")
877
878	    if not cache_path.exists():
879	        if rank == 0:
880	            print("Pre-computing token lengths (parallel, 32 workers)...")
881	        import multiprocessing as mp
882
883	        tok_dir = cfg.model.tokenizer.dir
884	        tok_model = str(Path(tok_dir) / "tokenizer.model")
885
886	        transcripts = list(train_dataset.transcripts)
887	        n_workers = min(32, mp.cpu_count())
888	        chunk_size = (len(transcripts) + n_workers - 1) // n_workers
889	        chunks = [transcripts[i:i + chunk_size] for i in range(0, len(transcripts), chunk_size)]
890
891	        if rank == 0:
892	            t_start = time.time()
893
894	        with mp.Pool(n_workers) as pool:
895	            results = pool.starmap(
896	                _tokenize_chunk_worker,
897	                [(chunk, tok_model) for chunk in chunks],
898	            )
899


Perfect! Now let me check how the model saves/loads to see if language embeddings are preserved:


550	            lang_bias = model._lang_embed(lang_ids)  # (B, 1, D)
551	            encoded = encoded + lang_bias
552	        return encoded, encoded_len
553
554	    model.encoder.forward = encoder_forward_with_lang
555	    return model
556
557
558	# ---------------------------------------------------------------------------
559	# Custom training step with language conditioning
560	# ---------------------------------------------------------------------------
561	def patch_training_step(model):
562	    """Patch the model's training_step to extract lang_ids from batch
563	    and pass them through to the encoder via language conditioning.
564	    """
565	    original_training_step = model.training_step
566
567	    def patched_training_step(batch, batch_idx):
568	        # NeMo expects batch = (signal, signal_len, tokens, token_len)
569	        # We added lang_ids as 5th element
570	        if isinstance(batch, (list, tuple)) and len(batch) >= 5:
571	            signal, signal_len, tokens, token_len, lang_ids = batch[:5]
572	            # Store lang_ids on model for encoder hook to pick up
573	            model._current_lang_ids = lang_ids.to(signal.device)
574	            # Pass standard 4-tuple to NeMo
575	            nemo_batch = (signal, signal_len, tokens, token_len)
576	            return original_training_step(nemo_batch, batch_idx)
577	        return original_training_step(batch, batch_idx)
578
579	    model.training_step = patched_training_step
580
581
582	def patch_encoder_for_lang(model):
583	    """Patch encoder forward to inject language IDs stored on model."""
584	    original_encoder_forward = model.encoder.forward
585
586	    def encoder_forward_with_lang(audio_signal, length):
587	        encoded, encoded_len = original_encoder_forward(audio_signal=audio_signal, length=length)
588	        if (hasattr(model, '_current_lang_ids') and hasattr(model, '_lang_embed')
589	                and model._current_lang_ids is not None):
590	            lang_ids = model._current_lang_ids
591	            # Safety: only apply if batch sizes match (skip during validation)
592	            if lang_ids.shape[0] == encoded.shape[0]:
593	                lang_bias = model._lang_embed(lang_ids)  # (B, D, 1)
594	                encoded = encoded + lang_bias
595	            # Clear after use to prevent stale values during validation
596	            model._current_lang_ids = None
597	        return encoded, encoded_len
598
599	    model.encoder.forward = encoder_forward_with_lang
600
601
602	# ---------------------------------------------------------------------------
603	# Main
604	# ---------------------------------------------------------------------------
605	def _tokenize_chunk_worker(texts, tokenizer_model_path):
606	    """Tokenize a chunk of texts in a worker process. Module-level for pickling."""
607	    import sentencepiece as spm
608	    sp = spm.SentencePieceProcessor()
609	    sp.Load(tokenizer_model_path)
610	    return [len(sp.EncodeAsIds(str(t).strip() or "<blank>")) for t in texts]
611
612
613	def main():
614	    parser = argparse.ArgumentParser(
615	        description="Production NeMo ASR training with TDT + differential LR"
616	    )
617	    parser.add_argument("--config", type=str, required=True)
618	    parser.add_argument("--max-steps", type=int, default=None)
619	    parser.add_argument("--devices", type=int, default=None)
620	    parser.add_argument("--train-parquet", type=str,
621	                        default="artifacts/phase3/production_train_final.parquet")
622	    parser.add_argument("--val-manifest", type=str,
623	                        default="data/manifests/stage1_prod_val_v2.jsonl")
624	    parser.add_argument("--resume-from-checkpoint", type=str, default=None)
625	    parser.add_argument("--resume-weights-only", action="store_true", default=False,
626	                        help="Resume model weights from checkpoint but reset optimizer (for config changes)")
627	    parser.add_argument("--pretrained-encoder", type=str, default=None,
628	                        help="Path to pretrained encoder weights (.pt file)")
629	    parser.add_argument("--smoke", action="store_true", default=False)
630	    parser.add_argument("--val-every", type=int, default=None)
631	    parser.add_argument("--log-every", type=int, default=None)
632
633	    # Data strategy args
634	    parser.add_argument("--max-batch-dur", type=float, default=120.0,
635	                        help="Max audio seconds per micro-batch (default: 120s)")
636	    parser.add_argument("--max-batch-size", type=int, default=16,
637	                        help="Global hard cap on samples per micro-batch")
638	    parser.add_argument("--max-tokens-in-batch", type=int, default=400,
639	                        help="Max total tokens per batch (controls joint tensor T_dec)")
640	    parser.add_argument("--temperature", type=float, default=0.3)
641	    parser.add_argument("--grad-accum", type=int, default=None)
642	    parser.add_argument("--num-workers", type=int, default=None)
643
644	    # v2 optimization args
645	    parser.add_argument("--encoder-lr", type=float, default=2e-5,
646	                        help="Learning rate for pretrained encoder")
647	    parser.add_argument("--head-lr", type=float, default=1e-3,
648	                        help="Learning rate for decoder/joint/CTC heads")
649	    parser.add_argument("--freeze-encoder-steps", type=int, default=5000,


Now let me create a summary document based on what I've found. Let me also check if there are any validation/inference hooks:


11-Optimizations v2:
12-  - TDT loss (matching pretrained Parakeet encoder objective)
13:  - Differential LR: encoder=2e-5, decoder/joint/CTC=1e-3
14-  - Encoder frozen for first 5K steps, then unfrozen
15-  - CTC loss warmup: 0→0.3 over first 3K steps
--
63-        torch.nn.init.normal_(self.embed.weight, mean=0.0, std=0.02)
64-
65:    def forward(self, lang_ids: torch.Tensor) -> torch.Tensor:
66-        """
67-        Args:
--
265-
266-class DifferentialLRCallback(pl.Callback):
267:    """Sets differential learning rates for encoder vs decoder/joint/CTC param groups.
268-
269-    After NeMo constructs the optimizer (which uses a single LR), this callback
--
308-        head_group['params'] = head_params
309-        head_group['lr'] = self.head_lr
310:        head_group['name'] = 'decoder_joint_ctc'
311-
312-        optimizer.param_groups.append(encoder_group)
--
375-    # Joint tensor size per sample in GB (bf16 = 2 bytes)
376-    joint_per_sample_gb = (enc_frames * max_t_dec * output_dim * 2) / 1e9
377:    # With gradient: 2x (forward + backward)
378-    joint_per_sample_gb *= 2
379-
--

[Showing results with pagination = limit: 30]


Perfect! Now I have enough information. Let me compile a comprehensive summary of the inference approach.

## Summary

I've thoroughly explored the Maya ASR Parakeet TDT 1.1B model codebase. Here's what I found about loading and running the model with language tag/embedding support during inference:

### Key Files Found

1. **Main Training Script**: `/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/scripts/train_prod.py` (13.5K lines)
   - Contains `LanguageEmbedding` class implementation
   - Shows how language conditioning is patched into the encoder
   - Demonstrates the full language embedding pipeline

2. **Model Config**: `/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/configs/train/stage1_prod_8xh200.yaml`
   - FastConformer encoder: 42 layers, 1024-dim hidden
   - TDT decoder + CTC head configuration
   - Model parameters: ~1,144M total (1,062M encoder + 81M heads)

3. **Documentation**: `/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/README.md`
   - Section "## Inference" provides NeMo API usage
   - Explains TDT-specific decoding configuration

### Language Support Architecture

**Language ID Mapping** (from train_prod.py, lines 49-53):
```
LANG_TO_ID = {
    "hi": 0, "bn": 1, "ta": 2, "te": 3, "mr": 4, "gu": 5,
    "kn": 6, "ml": 7, "pa": 8, "or": 9, "as": 10, "en": 11,
}
NUM_LANGUAGES = 12
```

**Language Embedding Module** (train_prod.py, lines 56-73):
- Takes language ID integers and projects to 1024-dim additive bias
- Initialized with `mean=0.0, std=0.02` to avoid disrupting pretrained encoder
- Output shape: `(B, 1024, 1)` — broadcast-addable to encoder output `(B, 1024, T)`
- Applied after encoder processes audio: `encoded = encoded + lang_bias`

### How to Use Language Tags at Inference

From the README (lines 134-175), the inference pattern is:

```python
import nemo.collections.asr as nemo_asr

# Load the trained model
model = nemo_asr.models.EncDecHybridRNNTCTCBPEModel.restore_from("path/to/maya_asr.nemo")
model.eval()
model.cuda()

# Transcribe audio files (standard NeMo API)
transcriptions = model.transcribe(["audio1.wav", "audio2.wav"])
print(transcriptions)
```

### Critical Implementation Details

**Language Embedding Forward Hook** (train_prod.py, lines 535-554):
The encoder forward is patched to inject language embeddings:
```python
def encoder_forward_with_lang(audio_signal, length, lang_ids=None):
    encoded, encoded_len = original_forward(audio_signal=audio_signal, length=length)
    if lang_ids is not None and hasattr(model, '_lang_embed'):
        lang_bias = model._lang_embed(lang_ids)  # (B, 1024, 1)
        encoded = encoded + lang_bias
    return encoded, encoded_len
```

**Training Step Modification** (train_prod.py, lines 561-577):
- Batch extended from 4-tuple to 5-tuple: `(signal, signal_len, tokens, token_len, lang_ids)`
- Language IDs stored on model and passed to encoder hook
- Standard NeMo 4-tuple passed to training step

**Encoder Patching for Inference** (train_prod.py, lines 582-599):
```python
def encoder_forward_with_lang(audio_signal, length):
    encoded, encoded_len = original_encoder_forward(...)
    if hasattr(model, '_current_lang_ids') and hasattr(model, '_lang_embed'):
        lang_ids = model._current_lang_ids
        if lang_ids.shape[0] == encoded.shape[0]:  # safety check
            lang_bias = model._lang_embed(lang_ids)  # (B, 1024, 1)
            encoded = encoded + lang_bias
        model._current_lang_ids = None  # clear after use
    return encoded, encoded_len
```

**Registration** (train_prod.py, lines 802-811):
- Language embedding registered as submodule: `model.register_module('_lang_embed_module', lang_embed)`
- This ensures it's saved with the checkpoint and moved to GPU automatically
- 12,288 parameters total for 12 languages × 1024-dim embeddings

### Inference Workflow

To use language conditioning during inference, you would need to:

1. Load the model: `model = nemo_asr.models.EncDecHybridRNNTCTCBPEModel.restore_from("maya_asr.nemo")`
2. Ensure language embedding module is loaded (saved automatically if checkpoint includes it)
3. Set language ID before transcription: `model._current_lang_ids = torch.tensor([LANG_TO_ID["hi"]])` for Hindi
4. Call standard NeMo transcribe: `model.transcribe(["audio.wav"])`

**Note**: The current codebase shows the architecture and training implementation but doesn't include a standalone inference script. The patching mechanism works during training; for inference, you'd need to re-patch the encoder forward method or extract/wrap the language embedding at inference time.

### Supported Decoding Modes

From config (lines 93-98):
- **TDT Greedy**: Primary (duration-aware, `durations=[0,1,2,3,4]`)
- **CTC**: Fallback (faster but less accurate)
- Both configured in model checkpoint