1	"""Tar offset reader: zero-extraction audio loading via positional I/O.
2
3	Reads FLAC bytes directly from tar files using os.pread() with byte offsets.
4	No tar extraction, no temp files, no runtime resampling.
5
6	Usage:
7	    reader = TarOffsetReader(max_fd_cache=256)
8	    waveform, sr = reader.read(tar_path, offset, nbytes)
9	    assert sr == 16000
10	"""
11
12	import io
13	import os
14	import threading
15	from collections import OrderedDict
16	from pathlib import Path
17
18	import soundfile
19
20
21	class TarOffsetReader:
22	    """Read audio from tar files via byte offsets using pread().
23
24	    Thread-safe with per-thread fd cache.
25	    """
26
27	    def __init__(self, max_fd_cache: int = 256, max_retries: int = 2):
28	        self.max_fd_cache = max_fd_cache
29	        self.max_retries = max_retries
30	        self._local = threading.local()
31
32	    def _get_cache(self) -> OrderedDict:
33	        if not hasattr(self._local, "fd_cache"):
34	            self._local.fd_cache = OrderedDict()
35	        return self._local.fd_cache
36
37	    def _get_fd(self, tar_path: str) -> int:
38	        cache = self._get_cache()
39	        if tar_path in cache:
40	            cache.move_to_end(tar_path)
41	            return cache[tar_path]
42	        # Evict oldest if cache full
43	        while len(cache) >= self.max_fd_cache:
44	            _, old_fd = cache.popitem(last=False)
45	            os.close(old_fd)
46	        fd = os.open(tar_path, os.O_RDONLY)
47	        cache[tar_path] = fd
48	        return fd
49
50	    def read(self, tar_path: str, offset: int, nbytes: int) -> tuple:
51	        """Read and decode audio from tar at given offset.
52
53	        Returns: (waveform_1d_numpy, sample_rate)
54	        Raises: RuntimeError on decode failure after retries.
55	        """
56	        last_err = None
57	        for attempt in range(self.max_retries + 1):
58	            try:
59	                fd = self._get_fd(tar_path)
60	                raw = os.pread(fd, nbytes, offset)
61	                if len(raw) != nbytes:
62	                    raise IOError(
63	                        f"Short read: got {len(raw)} expected {nbytes} "
64	                        f"at offset {offset} in {tar_path}"
65	                    )
66	                data, sr = soundfile.read(io.BytesIO(raw))
67	                if sr != 16000:
68	                    raise ValueError(f"Expected 16kHz, got {sr}Hz in {tar_path}")
69	                if len(data.shape) > 1:
70	                    data = data[:, 0]  # Take first channel if stereo
71	                return data, sr
72	            except Exception as e:
73	                last_err = e
74	                # Invalidate fd on error (might be stale)
75	                cache = self._get_cache()
76	                if tar_path in cache:
77	                    try:
78	                        os.close(cache.pop(tar_path))
79	                    except OSError:
80	                        pass
81	        raise RuntimeError(
82	            f"Failed to read after {self.max_retries + 1} attempts: "
83	            f"{tar_path}@{offset}+{nbytes}: {last_err}"
84	        )
85
86	    def close(self):
87	        """Close all cached file descriptors."""
88	        cache = self._get_cache()
89	        for fd in cache.values():
90	            try:
91	                os.close(fd)
92	            except OSError:
93	                pass
94	        cache.clear()
95
96	    def __del__(self):
97	        try:
98	            self.close()
99	        except Exception:
100	            pass


Bash

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


-rw-rw-r-- 1 ubuntu ubuntu  739 Mar 25 02:16 /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/configs/data/local_smoke.yaml
-rw-rw-r-- 1 ubuntu ubuntu 1.3M Mar 25 02:16 /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/configs/data/stage1_prod_input_cfg_v2.yaml
-rw-rw-r-- 1 ubuntu ubuntu 3.5K Mar 25 02:16 /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/configs/data/stage1_prod_input_cfg.yaml
-rw-rw-r-- 1 ubuntu ubuntu  579 Mar 25 02:16 /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/configs/data/stage1_smoke_input_cfg.yaml
-rw-rw-r-- 1 ubuntu ubuntu 4.5K Mar 27 05:10 /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/configs/train/stage1_prod_8xh200.yaml
-rw-rw-r-- 1 ubuntu ubuntu 4.2K Mar 25 02:16 /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/configs/train/stage1_prod_smoke_tarred.yaml
-rw-rw-r-- 1 ubuntu ubuntu 3.0K Mar 25 02:16 /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/configs/train/stage1_prod_smoke_test.yaml
-rw-rw-r-- 1 ubuntu ubuntu 3.3K Mar 25 02:16 /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/configs/train/stage1_smoke.yaml


Now let me examine scripts and check for any custom model layer additions:


find /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1 -type f \( -name "*.py" -o -name "*.yaml" \) | grep -v __pycache__ | grep -v ".git" | head -60


Now let me check some of the other Python scripts to understand data pipeline:


145	    def _pack_bucket_vectorized(self, indices: np.ndarray, bucket_bs: int,
146	                                 rng_state: np.random.RandomState) -> list[list[int]]:
147	        """Pack a single bucket into batches using vectorized operations.
148
149	        Uses numpy for the weighted permutation and greedy packing.
150	        """
151	        n = len(indices)
152	        if n == 0:
153	            return []
154
155	        # Weighted random permutation: U^(1/w) trick, fully vectorized
156	        weights = self._sample_weights[indices]
157	        weights = np.maximum(weights, 1e-12)
158	        u = rng_state.random(n)
159	        keys = np.power(u, 1.0 / weights)
160
161	        # Sort descending by key
162	        order = np.argsort(-keys)
163	        sorted_indices = indices[order]
164
165	        # Get durations and token lengths for sorted samples
166	        durs = self.durations[sorted_indices]
167	        toks = self.token_lengths[sorted_indices]
168
169	        # Greedy packing with triple cap (vectorized cumsum approach)
170	        # We compute cumulative duration and tokens, then find split points
171	        cum_dur = np.cumsum(durs)
172	        cum_tok = np.cumsum(toks)
173
174	        batches = []
175	        start = 0
176	        while start < n:
177	            # Find how many samples we can fit from start position
178	            remaining = n - start
179	            max_count = min(remaining, bucket_bs)
180
181	            if max_count == 0:
182	                break
183
184	            # Binary search for the largest count that fits all three caps
185	            # Durations: cum_dur[start+count-1] - cum_dur[start-1] <= max_batch_duration
186	            # Tokens: cum_tok[start+count-1] - cum_tok[start-1] <= max_tokens_in_batch
187	            base_dur = cum_dur[start - 1] if start > 0 else 0.0
188	            base_tok = cum_tok[start - 1] if start > 0 else 0
189
190	            # Vectorized: find all valid positions
191	            end_range = min(start + max_count, n)
192	            batch_durs = cum_dur[start:end_range] - base_dur
193	            batch_toks = cum_tok[start:end_range] - base_tok
194
195	            # Find first position that exceeds either cap
196	            dur_ok = batch_durs <= self.max_batch_duration
197	            tok_ok = batch_toks <= self.max_tokens_in_batch
198
199	            valid = dur_ok & tok_ok
200	            if valid.all():
201	                count = len(valid)
202	            elif not valid[0]:
203	                # Even a single sample exceeds a cap — take it anyway (BS=1)
204	                count = 1
205	            else:
206	                # Find first False
207	                count = int(np.argmin(valid))
208	                if count == 0:
209	                    count = 1
210
211	            batch = sorted_indices[start:start + count].tolist()
212	            if batch and (not self.drop_last or len(batch) >= 2):
213	                batches.append(batch)
214
215	            start += count
216
217	        return batches
218
219	    def __iter__(self) -> Iterator[list[int]]:
220	        rank = self.rank
221	        is_main = rank == 0
222	        t0 = time.time()
223
224	        # Use numpy RNG for vectorized operations
225	        rng_np = np.random.RandomState(self.seed + self.epoch)
226
227	        # Step 1: Group samples by bucket (vectorized)
228	        unique_buckets = np.unique(self._bucket_ids)
229	        bucket_indices = {
230	            bid: np.where(self._bucket_ids == bid)[0]
231	            for bid in unique_buckets
232	        }
233
234	        if is_main:
235	            t1 = time.time()
236	            print(f"  [Sampler] Bucketing: {t1-t0:.1f}s")
237
238	        # Step 2+3: Weighted permutation + greedy packing per bucket (vectorized)
239	        all_batches = []
240	        bucket_stats = {}
241
242	        for bucket_id in unique_buckets:
243	            indices = bucket_indices[bucket_id]
244	            bucket_bs = self._get_bucket_bs(int(bucket_id))
245
246	            batches = self._pack_bucket_vectorized(indices, bucket_bs, rng_np)
247	            all_batches.extend(batches)
248
249	            lo = self.boundaries[bucket_id - 1] if bucket_id > 0 else 0.0
250	            hi = self.boundaries[bucket_id] if bucket_id < len(self.boundaries) else 99.0
251	            bucket_stats[int(bucket_id)] = {
252	                "range": f"{lo:.0f}-{hi:.0f}s",
253	                "samples": len(indices),
254	                "batches": len(batches),
255	                "bs_cap": bucket_bs,
256	            }
257
258	        if is_main:
259	            t2 = time.time()
260	            print(f"  [Sampler] Packing: {t2-t1:.1f}s ({len(all_batches):,} batches)")
261
262	        # Step 4: Sort by compute cost (vectorized batch cost)
263	        # Pre-compute max duration per batch using numpy
264	        batch_costs = np.array([
265	            len(b) * self.durations[b].max() for b in all_batches
266	        ], dtype=np.float32)
267	        cost_order = np.argsort(-batch_costs)
268	        all_batches = [all_batches[i] for i in cost_order]
269
270	        if is_main:
271	            t3 = time.time()
272	            print(f"  [Sampler] Cost sort: {t3-t2:.1f}s")
273
274	        # Step 5: DDP sharding
275	        usable = (len(all_batches) // self.world_size) * self.world_size
276	        all_batches = all_batches[:usable]
277	        rank_batches = all_batches[rank::self.world_size]
278
279	        # Step 6: Grad accum alignment
280	        if self.grad_accum > 1:
281	            trim = len(rank_batches) - (len(rank_batches) % self.grad_accum)
282	            rank_batches = rank_batches[:trim]
283
284	        # Step 7: Shuffle
285	        if self.shuffle:
286	            rng_np.shuffle(rank_batches)
287
288	        if is_main:
289	            t4 = time.time()
290	            print(f"  [Sampler] Shard+shuffle: {t4-t3:.1f}s")
291
292	        # Collect stats (vectorized)
293	        if rank_batches:
294	            batch_sizes = np.array([len(b) for b in rank_batches])


total 4.0K
drwxrwxr-x 2 ubuntu ubuntu 4.0K Mar 27 05:26 stage1_prod_bpe


Let me check the tokenizer directory and look at the data configs:


total 1.5M
-rw-rw-r-- 1 ubuntu ubuntu 962K Mar 27 05:26 tokenizer.model
-rw-rw-r-- 1 ubuntu ubuntu 472K Mar 27 05:26 vocab.txt


1	# Smoke-run config for local data ingestion
2	# Processes 1 shard per language for en+hi only
3
4	build_manifest:
5	  output: data/manifests/smoke.jsonl
6	  languages: [en, hi]
7	  max_shards: 1
8	  sources: [final-export]
9	  min_quality: 0.0
10
11	split_manifest:
12	  input: data/manifests/smoke.jsonl
13	  train_output: data/manifests/smoke_train.jsonl
14	  val_output: data/manifests/smoke_val.jsonl
15	  val_ratio: 0.01
16	  seed: 42
17	  stratify_by_lang: true
18
19	generate_input_cfg:
20	  manifests: [data/manifests/smoke_train.jsonl]
21	  output: configs/data/stage1_smoke_input_cfg.yaml
22
23	validate_manifest:
24	  min_duration: 0.1
25	  max_duration: 60.0
26	  required_fields:
27	    - audio_filepath
28	    - text
29	    - duration
30	    - lang
31	    - taskname
32	    - source_lang
33	    - target_lang
34


Now let me look at a few more key data processing scripts:


1	#!/usr/bin/env python3
2	"""Build NeMo JSONL manifests from local sft_data shard metadata.
3
4	Reads metadata.parquet + audio_index.parquet from each shard and produces
5	NeMo-compatible JSONL manifest lines:
6	  {"audio_filepath", "text", "duration", "lang", "taskname", "source_lang", "target_lang"}
7
8	Audio filepaths point into the shard's audio.tar using pipe: notation
9	(tar_member_name within the tar), suitable for Lhotse/NeMo tarred datasets.
10
11	Usage:
12	  # Smoke (default): 1 shard per language, en+hi only
13	  python scripts/build_manifest.py --output data/manifests/smoke.jsonl
14
15	  # Custom smoke
16	  python scripts/build_manifest.py --languages en hi ta --max-shards 2 \
17	    --output data/manifests/custom.jsonl
18
19	  # Full build (all languages, all shards)
20	  python scripts/build_manifest.py --max-shards 0 --output data/manifests/full.jsonl
21	"""
22
23	import argparse
24	import json
25	import sys
26	from pathlib import Path
27
28	import pyarrow.parquet as pq
29	from tqdm import tqdm
30
31	# Add src to path
32	sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
33	from maya_asr.config import FINAL_EXPORT_ROOT, LANGUAGES
34
35	# All data source roots that follow the same shard format
36	DATA_SOURCES = {
37	    "final-export": FINAL_EXPORT_ROOT,
38	}
39
40	# Additional sources with lang=XX directory structure at top level
41	LANG_DIR_SOURCES = {
42	    "indicvoices": Path("/root/sft_data/indicvoices"),
43	    "indicvoices-r": Path("/root/sft_data/indicvoices-r"),
44	    "josh": Path("/root/sft_data/josh"),
45	    "joshdelivery": Path("/root/sft_data/joshdelivery"),
46	}
47
48	# Sources with lang=XX subdirectory structure
49	HIVE_SOURCES = {
50	    "ears": Path("/root/sft_data/ears"),
51	    "expresso": Path("/root/sft_data/expresso"),
52	    "globe": Path("/root/sft_data/globe"),
53	    "librittsr": Path("/root/sft_data/librittsr"),
54	    "ljspeech": Path("/root/sft_data/ljspeech"),
55	    "vctk": Path("/root/sft_data/vctk"),
56	}
57
58
59	def discover_shards(source_root: Path, languages: list[str]) -> list[tuple[str, Path]]:
60	    """Discover shard directories for given languages under a source root.
61
62	    Handles two directory layouts:
63	    1. Hive-partitioned: source_root/lang={code}/{shard_id}/
64	    2. Plain language dirs: source_root/{code}/{shard_id}/
65	    """
66	    shards = []
67	    for lang in languages:
68	        # Try hive-partitioned first
69	        lang_dir = source_root / f"lang={lang}"
70	        if not lang_dir.exists():
71	            # Try plain language directory
72	            lang_dir = source_root / lang
73	        if not lang_dir.exists():
74	            continue
75	        for shard_dir in sorted(lang_dir.iterdir()):
76	            if shard_dir.is_dir() and (shard_dir / "metadata.parquet").exists():
77	                shards.append((lang, shard_dir))
78	    return shards
79
80
81	def process_shard(lang: str, shard_dir: Path, min_quality: float = 0.0) -> tuple[list[dict], int]:
82	    """Convert one shard's metadata into NeMo manifest rows.
83
84	    Returns (rows, dropped_low_quality_count).
85	    """
86	    meta_path = shard_dir / "metadata.parquet"
87	    audio_index_path = shard_dir / "audio_index.parquet"
88	    audio_tar_path = shard_dir / "audio.tar"
89
90	    # Read metadata
91	    meta_cols = [
92	        "segment_id",
93	        "transcription_native",
94	        "duration_s",
95	        "lang",
96	        "final_bucket",
97	        "tx_quality_score",
98	    ]
99	    try:
100	        meta_table = pq.read_table(meta_path, columns=meta_cols)
101	    except Exception as e:
102	        print(f"  WARN: Cannot read {meta_path}: {e}", file=sys.stderr)
103	        return [], 0
104
105	    meta_df = meta_table.to_pandas()
106
107	    # Read audio index for tar member mapping
108	    audio_lookup = {}
109	    if audio_index_path.exists():
110	        try:
111	            ai_table = pq.read_table(audio_index_path, columns=["segment_id", "tar_member_name"])
112	            ai_df = ai_table.to_pandas()
113	            audio_lookup = dict(zip(ai_df["segment_id"], ai_df["tar_member_name"]))
114	        except Exception as e:
115	            print(f"  WARN: Cannot read {audio_index_path}: {e}", file=sys.stderr)
116
117	    rows = []
118	    dropped_quality = 0
119	    for _, row in meta_df.iterrows():
120	        text = row.get("transcription_native", "")
121	        if not isinstance(text, str) or not text.strip():
122	            continue
123
124	        duration = row.get("duration_s")
125	        if duration is None or duration <= 0 or duration > 60:
126	            continue
127
128	        # Apply quality filter
129	        quality = row.get("tx_quality_score")
130	        if min_quality > 0 and (quality is None or float(quality) < min_quality):
131	            dropped_quality += 1
132	            continue
133
134	        segment_id = row["segment_id"]
135	        tar_member = audio_lookup.get(segment_id, segment_id)
136	        row_lang = str(row.get("lang", lang))
137
138	        rows.append(
139	            {
140	                "audio_filepath": str(audio_tar_path),
141	                "tar_member": tar_member,
142	                "text": text.strip(),
143	                "duration": round(float(duration), 3),
144	                "lang": row_lang,
145	                "taskname": "asr",
146	                "source_lang": row_lang,
147	                "target_lang": row_lang,
148	            }
149	        )
150


10-
11-Usage:
12:  # Smoke (default): 1 shard per language, en+hi only
13-  python scripts/build_manifest.py --output data/manifests/smoke.jsonl
14-
15-  # Custom smoke
16:  python scripts/build_manifest.py --languages en hi ta --max-shards 2 \
17-    --output data/manifests/custom.jsonl
18-
19:  # Full build (all languages, all shards)
20-  python scripts/build_manifest.py --max-shards 0 --output data/manifests/full.jsonl
21-"""
--
31-# Add src to path
32-sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
33:from maya_asr.config import FINAL_EXPORT_ROOT, LANGUAGES
34-
35-# All data source roots that follow the same shard format
--
39-
40-# Additional sources with lang=XX directory structure at top level
41:LANG_DIR_SOURCES = {
42-    "indicvoices": Path("/root/sft_data/indicvoices"),
43-    "indicvoices-r": Path("/root/sft_data/indicvoices-r"),
--
57-
58-
59:def discover_shards(source_root: Path, languages: list[str]) -> list[tuple[str, Path]]:
60:    """Discover shard directories for given languages under a source root.
61-
62-    Handles two directory layouts:
63-    1. Hive-partitioned: source_root/lang={code}/{shard_id}/
64:    2. Plain language dirs: source_root/{code}/{shard_id}/
65-    """
66-    shards = []
67:    for lang in languages:
68-        # Try hive-partitioned first
69-        lang_dir = source_root / f"lang={lang}"
70-        if not lang_dir.exists():
71:            # Try plain language directory
72-            lang_dir = source_root / lang
73-        if not lang_dir.exists():
--
161-    )
162-    parser.add_argument(
163:        "--languages",
164-        nargs="+",
165-        default=["en", "hi"],
166:        choices=LANGUAGES,
167:        help="Languages to process (default: en hi)",
168-    )
169-    parser.add_argument(
--
171-        type=int,
172-        default=1,
173:        help="Max shards per language per source (0=all, default=1 for smoke)",
174-    )
175-    parser.add_argument(
--
201-        if src in DATA_SOURCES:
202-            source_roots[src] = DATA_SOURCES[src]
203:        elif src in LANG_DIR_SOURCES:
204:            source_roots[src] = LANG_DIR_SOURCES[src]
205-        elif src in HIVE_SOURCES:
206-            source_roots[src] = HIVE_SOURCES[src]
--
209-
210-    # Discover and process shards
211:    stats = {lang: {"segments": 0, "duration_h": 0.0} for lang in args.languages}
212-    total_rows = 0
213-    total_dropped_quality = 0
--
219-        for source_name, source_root in source_roots.items():
220-            print(f"\n=== Source: {source_name} ({source_root}) ===")
221:            shards = discover_shards(source_root, args.languages)
222-
223-            if args.max_shards > 0:
224:                # Limit per language
225-                limited = []
226-                lang_counts = {}
--
259-    if total_dropped_dupes > 0:
260-        print(f"Dropped (duplicates):  {total_dropped_dupes:,}")
261:    print("\nPer-language breakdown:")
262-    print(f"  {'Lang':<6} {'Segments':>10} {'Hours':>10}")
263-    print(f"  {'-' * 6} {'-' * 10} {'-' * 10}")


Perfect! Let me gather more information about the architecture and check checkpoint handling:


1	#!/usr/bin/env python3
2	"""Restore a training checkpoint from Cloudflare R2.
3
4	Downloads .ckpt/.nemo files and verifies integrity via SHA-256.
5
6	Usage:
7	  # Restore latest step
8	  python3 scripts/restore_checkpoint.py \
9	    --model-name maya-asr-stage1 --target-dir /tmp/restored
10
11	  # Restore specific step
12	  python3 scripts/restore_checkpoint.py \
13	    --model-name maya-asr-stage1 --step 50 --target-dir /tmp/restored
14
15	  # List available steps
16	  python3 scripts/restore_checkpoint.py \
17	    --model-name maya-asr-stage1 --list
18	"""
19
20	import argparse
21	import hashlib
22	import json
23	import os
24	import re
25	import sys
26	from pathlib import Path
27
28	import boto3
29	from dotenv import load_dotenv
30
31	ENV_FILE = Path(__file__).resolve().parent.parent / ".env"
32	if ENV_FILE.exists():
33	    load_dotenv(ENV_FILE)
34
35
36	def get_s3_client():
37	    endpoint = os.environ.get("R2_ENDPOINT_URL")
38	    key_id = os.environ.get("R2_ACCESS_KEY_ID")
39	    secret = os.environ.get("R2_SECRET_ACCESS_KEY")
40	    if not all([endpoint, key_id, secret]):
41	        print("ERROR: R2 credentials not set.", file=sys.stderr)
42	        sys.exit(1)
43	    return boto3.client(
44	        "s3",
45	        endpoint_url=endpoint,
46	        aws_access_key_id=key_id,
47	        aws_secret_access_key=secret,
48	        region_name="auto",
49	    )
50
51
52	def list_steps(s3, bucket: str, model_name: str) -> list[int]:
53	    """List all available step numbers for a model on R2."""
54	    prefix = f"{model_name}/"
55	    paginator = s3.get_paginator("list_objects_v2")
56	    steps = set()
57	    for page in paginator.paginate(Bucket=bucket, Prefix=prefix, Delimiter="/"):
58	        for cp in page.get("CommonPrefixes", []):
59	            # cp["Prefix"] = "model/step_00000050/"
60	            m = re.search(r"step_(\d+)", cp["Prefix"])
61	            if m:
62	                steps.add(int(m.group(1)))
63	    return sorted(steps)
64
65
66	def _sha256(path: Path) -> str:
67	    h = hashlib.sha256()
68	    with open(path, "rb") as f:
69	        for chunk in iter(lambda: f.read(1 << 20), b""):
70	            h.update(chunk)
71	    return h.hexdigest()
72
73
74	def main():
75	    parser = argparse.ArgumentParser(description="Restore checkpoint from R2")
76	    parser.add_argument("--model-name", type=str, required=True)
77	    parser.add_argument("--step", type=int, default=None, help="Step to restore (default: latest)")
78	    parser.add_argument("--target-dir", type=Path, default=None)
79	    parser.add_argument("--bucket", type=str, default=None)
80	    parser.add_argument("--list", action="store_true", default=False, help="List available steps")
81	    args = parser.parse_args()
82
83	    bucket = args.bucket or os.environ.get("R2_BUCKET_CHECKPOINTS", "ptcheckpoints")
84	    s3 = get_s3_client()
85
86	    # List mode
87	    if args.list:
88	        steps = list_steps(s3, bucket, args.model_name)
89	        if not steps:
90	            print(f"No checkpoints found for {args.model_name} in {bucket}")
91	            sys.exit(1)
92	        print(f"Available steps for {args.model_name}:")
93	        for step in steps:
94	            print(f"  step_{step:08d}")
95	        return
96
97	    # Determine step
98	    if args.step is not None:
99	        step = args.step
100	    else:


grep -r "Gemma\|gemma\|projection\|LLM" /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1 --include="*.py" --include="*.md" --include="*.yaml" 2>/dev/null | head -20


/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/README.md:| **CTC Head** | Linear projection over encoder output | ~33M |
/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/README.md:### Stage 2: Gemma LLM Decoder (planned)
/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/README.md:Replaces the RNN decoder with a Gemma LLM for improved language modeling, particularly for low-resource Indic languages.
/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/lf_asr/__init__.py:"""Maya ASR — FastConformer + Gemma 3 1B Speech-to-Text."""
/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/src/maya_asr/__init__.py:"""Maya ASR — Multilingual Indic ASR with FastConformer + Gemma."""
/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/PLAN.md:# Maya ASR: FastConformer + Gemma Transcription Build Guide
/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/PLAN.md: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.
/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/PLAN.md:│  │  FastConformer Encoder │───►│Projection│───►│ Gemma 3      │   │
/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/PLAN.md:**Decoder: Gemma 3 (Small Variant)**
/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/PLAN.md:Use the smallest viable Gemma 3 variant as a decoder-only LLM. Candidates:
/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/PLAN.md:| `google/gemma-3-1b-pt` | 1B | 1536 | 26 | Smallest official Gemma 3 |
/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/PLAN.md:| `google/gemma-2-2b` | 2.6B | 2304 | 26 | Gemma 2, good Indic but larger |
/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/PLAN.md:**Primary choice: Gemma 3 1B** — it is already the smallest Gemma 3 release. Its 262K SentencePiece tokenizer has excellent Indic coverage:
/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/PLAN.md:| Language | Gemma 3 Fertility | Dedicated Tokens |
/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/PLAN.md: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.
/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/PLAN.md:Linear projection outperforms Q-Former for speech-LLM coupling ([MLC-SLM 2025](https://arxiv.org/html/2601.01461v1): Linear 11.05% WER vs Q-Former 11.51%).
/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/PLAN.md:This downsamples encoder output by 2x temporally (on top of 8x from FastConformer), then maps from encoder dimension (1024) to Gemma 3 1B's hidden_size (1536).
/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/PLAN.md:### 3.1 Decoder Tokenizer: Gemma 3's Tokenizer
/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/PLAN.md: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.
/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/PLAN.md:Fine-tune the encoder on all ~150K hours using hybrid CTC+RNNT loss. This creates a strong Indic speech feature extractor before coupling with the LLM decoder.


Excellent! Now let me check for any existing checkpoints and gather final details:


find /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1 -type f \( -name "*.ckpt" -o -name "*.pt" -o -name "*.bin" -o -name "*.nemo" \) 2>/dev/null


(Bash completed with no output)


ls -lh /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/pretrained/ 2>/dev/null || echo "No pretrained directory"


No pretrained directory


Now let me create a comprehensive summary document. Let me check a few more key files to ensure completeness:


"""PyTorch Dataset for tar-offset audio loading.

Reads audio via TarOffsetReader (pread, no extraction).
Returns waveform + transcript + metadata for NeMo training.
"""

import time

import numpy as np
import torch
from torch.utils.data import Dataset

from .tar_offset_reader import TarOffsetReader

class TarOffsetASRDataset(Dataset):
    """ASR dataset that reads audio from tar files via byte offsets.

    Args:
        manifest_df: DataFrame with columns:
            tar_path, tar_offset_data, tar_nbytes, duration_s, transcript, language
        max_duration: Maximum audio duration in seconds (skip longer).
        min_duration: Minimum audio duration in seconds (skip shorter).
    """

    def __init__(
        self,
        manifest_df,
        max_duration: float = 30.0,
        min_duration: float = 0.1,
        sample_rate: int = 16000,
    ):
        # Filter by duration
        mask = (manifest_df["duration_s"] >= min_duration) & (
            manifest_df["duration_s"] <= max_duration
        )
        self.df = manifest_df[mask].reset_index(drop=True)
        self.sample_rate = sample_rate
        self.reader = TarOffsetReader(max_fd_cache=512)

        # Pre-extract arrays for fast indexing
        self.tar_paths = self.df["tar_path"].values
        self.offsets = self.df["tar_offset_data"].values.astype(np.int64)
        self.nbytes = self.df["tar_nbytes"].values.astype(np.int64)
        self.transcripts = self.df["transcript"].values
        self.languages = self.df["language"].values
        self.durations = self.df["duration_s"].values.astype(np.float32)

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        t0 = time.monotonic()

        tar_path = str(self.tar_paths[idx])
        offset = int(self.offsets[idx])
        nbytes = int(self.nbytes[idx])

        waveform, sr = self.reader.read(tar_path, offset, nbytes)
        assert sr == self.sample_rate, f"Expected {self.sample_rate}, got {sr}"

        decode_time = time.monotonic() - t0

        return {
            "audio": torch.from_numpy(waveform).float(),
            "audio_len": len(waveform),
            "transcript": str(self.transcripts[idx]),
            "language": str(self.languages[idx]),
            "duration": float(self.durations[idx]),
            "decode_time_ms": decode_time * 1000,
        }

    def close(self):
        self.reader.close()

def collate_asr_batch(batch):
    """Collate ASR samples with padding."""
    max_len = max(b["audio_len"] for b in batch)

    audio_padded = torch.zeros(len(batch), max_len)
    audio_lens = torch.zeros(len(batch), dtype=torch.long)
    transcripts = []
    languages = []

    padding_samples = 0
    total_samples = 0

    for i, b in enumerate(batch):
        audio_padded[i, : b["audio_len"]] = b["audio"]
        audio_lens[i] = b["audio_len"]
        transcripts.append(b["transcript"])
        languages.append(b["language"])
        padding_samples += max_len - b["audio_len"]
        total_samples += max_len

    padding_ratio = padding_samples / max(total_samples, 1)

    return {
        "audio": audio_padded,


# Stage 1 Production Config: FastConformer Hybrid TDT+CTC on full dataset
# Hardware: 8x NVIDIA H200 (140 GB HBM3 each)
# Target: ~200K steps, ~5-7 days
#
# Pretrained encoder: Parakeet TDT 1.1B (42-layer FastConformer, English)
# Tokenizer: encoder-side 32K BPE (quality-filtered text, 11 Indic + English)
# Decoder/joint: randomly initialized for 32K vocab, then finetuned
#
# Optimizations v2:
#   - TDT loss (matching pretrained encoder objective) instead of RNNT
#   - Differential LR: encoder=2e-5, decoder/joint/CTC=1e-3
#   - Encoder frozen for first 5K steps, then unfrozen at low LR
#   - CTC loss warmup: 0 -> 0.3 over first 3K steps
#   - Language embedding conditioning (lang_id -> encoder)

--
    num_extra_outputs: 5  # TDT duration heads (durations=[0,1,2,3,4])
    fuse_loss_wer: true
    fused_batch_size: 2   # sub-batch joint computation → caps peak VRAM
    jointnet:
      joint_hidden: 640
      activation: relu
--
    loss_name: tdt
    tdt_kwargs:
      durations: [0, 1, 2, 3, 4]
      sigma: 0.02
      omega: 0.1
      fastemit_lambda: 0.0
      clamp: -1.0
--
    model_type: tdt
    durations: [0, 1, 2, 3, 4]
    greedy:
      max_symbols: 10

  aux_ctc:
--
  name: maya_asr_stage1_tdt_v2
  create_tensorboard_logger: true
  create_wandb_logger: true
  wandb_logger_kwargs:
    project: maya-asr
    name: tdt-parakeet-finetune-stage1-v2
    resume: allow
  resume_if_exists: false
  resume_ignore_no_checkpoint: true
  create_checkpoint_callback: true
  checkpoint_callback_params:


Perfect! Now I have comprehensive information. Let me create a detailed summary report:

## Comprehensive Analysis: Maya ASR Hybrid FastConformer-RNNT-Stage1

I've completed a thorough exploration of the Maya ASR project. Here's the complete breakdown:

### 1. DIRECTORY STRUCTURE & PROJECT OVERVIEW

**Root Directory:** `/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/`

**Key Directories:**
- `configs/` — Training and data YAML configurations
- `lf_asr/data/` — Custom data loading and sampling modules
- `src/maya_asr/` — Core library with configuration constants
- `scripts/` — Training, data pipeline, checkpoint management scripts
- `tools/` — Data preparation and validation utilities
- `tests/` — Unit and integration tests
- `tokenizers/stage1_prod_bpe/` — 32K BPE SentencePiece tokenizer (1.5MB model + vocab)

---

### 2. ARCHITECTURE OVERVIEW

**Model Type:** `EncDecHybridRNNTCTCBPEModel` (NeMo native)

**Architecture Diagram:**
```
Audio (16 kHz)
    ↓
Mel Spectrogram (80-dim, 16000 Hz sample rate)
    ↓
FastConformer Encoder (42 layers, 1024-dim, 8x subsampling)
    ├─ Pretrained from Parakeet TDT 1.1B encoder
    └─ + Language Embedding (12 languages → 1024-dim additive bias)
    ↓
    ├─────────────────┬─────────────────┐
    ↓                 ↓                 ↓
TDT Decoder      Joint Network      CTC Head
(2-layer RNN)    (640-dim hidden)    (Linear projection)
(640-dim)        + 5 duration heads  (~33M params)
(~7M params)     (~42M params)
    ↓                 ↓
    └─────────┬───────┘
              ↓
         [Primary + Auxiliary Loss]
         TDT (transducer) + CTC
```

**Total Parameters:** ~1,144M (1.1B)

**Component Breakdown:**
- **Encoder:** 42-layer FastConformer, 1024-dim hidden, 8x subsampling (dw_striding), ReLU activation
  - ~1,062M parameters
  - Pretrained from Parakeet TDT 1.1B (frozen first 5K steps, then fine-tuned at 2e-5 LR)

- **Decoder:** 2-layer RNN with prediction network
  - pred_hidden=640, pred_rnn_layers=2
  - ~7M parameters
  - Randomly initialized, trained at 1e-3 LR

- **Joint Network:** RNNTJoint with extra duration outputs
  - joint_hidden=640
  - num_extra_outputs=5 (TDT duration heads for durations [0,1,2,3,4])
  - ~42M parameters
  - Fused computation with fused_batch_size=2 to cap VRAM

- **CTC Head:** Linear projection over encoder output
  - ~33M parameters
  - Auxiliary loss weighted 0.3 (warmed up from 0→0.3 over 3K steps)

- **Language Embedding (NEW):** Per-language additive bias
  - 12 languages → 1024-dim embeddings
  - ~12K parameters
  - Initialized with std=0.02 to avoid disrupting pretrained weights

---

### 3. LANGUAGE CONDITIONING & INPUT

**Language as Model Input:**

1. **Language Embedding Module** (`scripts/train_prod.py`, lines 56-73):
```python
class LanguageEmbedding(torch.nn.Module):
    """Projects a language ID into a per-frame additive bias for the encoder."""
    def __init__(self, num_languages: int, embed_dim: int):
        super().__init__()
        self.embed = torch.nn.Embedding(num_languages, embed_dim)
        torch.nn.init.normal_(self.embed.weight, mean=0.0, std=0.02)

    def forward(self, lang_ids: torch.Tensor) -> torch.Tensor:
        # lang_ids: (B,) → returns (B, embed_dim, 1) for broadcasting
        return self.embed(lang_ids).unsqueeze(2)
```

2. **Language ID Mapping** (lines 49-53):
```python
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,
}
```

3. **Data Loading with Language IDs** (`scripts/train_prod.py`, lines 451-455):
```python
self.lang_ids = np.array(
    [LANG_TO_ID.get(str(lang).strip().lower(), LANG_TO_ID["en"])
     for lang in self.languages],
    dtype=np.int64,
)
```

4. **Application in Encoder** (`scripts/train_prod.py`, lines 582-599):
```python
def patch_encoder_for_lang(model):
    """Patch encoder forward to inject language IDs stored on model."""
    original_encoder_forward = model.encoder.forward

    def encoder_forward_with_lang(audio_signal, length):
        encoded, encoded_len = original_encoder_forward(audio_signal=audio_signal,
                                                        length=length)
        if (hasattr(model, '_current_lang_ids') and hasattr(model, '_lang_embed')
                and model._current_lang_ids is not None):
            lang_ids = model._current_lang_ids
            if lang_ids.shape[0] == encoded.shape[0]:
                lang_bias = model._lang_embed(lang_ids)  # (B, D, 1)
                encoded = encoded + lang_bias  # (B, D, T) + (B, D, 1)
        return encoded, encoded_len

    model.encoder.forward = encoder_forward_with_lang
```

**Supported Languages:** Hindi (hi), Bengali (bn), Tamil (ta), Telugu (te), Marathi (mr), Gujarati (gu), Kannada (kn), Malayalam (ml), Punjabi (pa), Odia (or), Assamese (as), English (en)

---

### 4. CUSTOM LAYERS & MODIFICATIONS

**LanguageEmbedding Module** (NEW custom layer):
- **File:** `/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/scripts/train_prod.py` (lines 56-73)
- **Purpose:** Injects language signal via per-frame additive bias
- **Architecture:** Embedding layer (12 → 1024-dim)
- **Initialization:** Normal(mean=0, std=0.02) to preserve pretrained encoder
- **Integration:** Hooked into encoder.forward via patch

**TDT (Temporal Duration Transducer) Modifications:**
- **Joint Network:** Added `num_extra_outputs: 5` for duration predictions (0, 1, 2, 3, 4 frames)
- **Loss:** Changed from RNNT to TDT loss with:
  - durations: [0, 1, 2, 3, 4]
  - sigma: 0.02 (softness parameter)
  - omega: 0.1 (regularization)
  - fastemit_lambda: 0.0

**Fused Joint Computation:**
- Uses `fuse_loss_wer: true` with `fused_batch_size: 2` to split joint tensor computation into sub-batches

---

### 5. TRAINING SCRIPTS & OPTIMIZATIONS

**Main Training Script:** `/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/scripts/train_prod.py` (1,023 lines)

**Key Callbacks & Optimizations:**

1. **EncoderFreezeCallback** (lines 162-201):
   - Freezes encoder for first 5,000 optimizer steps
   - Allows decoder/joint/CTC heads to warm up without corrupting pretrained weights
   - Unfreezes encoder at 2e-5 LR thereafter

2. **CTCWarmupCallback** (lines 245-263):
   - Ramps CTC loss weight from 0 → 0.3 over first 3,000 steps
   - Prevents random CTC gradients from disrupting early training

3. **DifferentialLRCallback** (lines 266-332):
   - **Encoder:** 2e-5 LR (pretrained, low learning rate)
   - **Decoder/Joint/CTC:** 1e-3 LR (random init, fast adaptation)
   - Splits optimizer param groups and syncs with LR scheduler

4. **ThroughputCallback** (lines 79-159):
   - Logs VRAM, samples/s, audio-s/s, RTF/GPU, MFU every N steps
   - Per-param-group LR logging to wandb/tensorboard

5. **SamplerStateCallback** (lines 204-242):
   - Saves sampler state for resumption without reprocessing data

**Training Data Strategy:**

- **ProductionBatchSampler** (`/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/lf_asr/data/production_sampler.py`):
  - **Duration Bucketing:** Groups samples into 5 buckets ([0-4], [4-8], [8-12], [12-16], [16-20]s)
  - **Language Temperature Rebalancing (T=0.3):** Upweights underrepresented languages
  - **Triple-Cap OOM Protection:**
    1. Per-bucket BS caps (24→6 samples depending on bucket)
    2. max_batch_duration (120s audio per micro-batch)
    3. max_tokens_in_batch (400 tokens, controls T_dec in joint tensor)
  - **Compute-Cost DDP Sharding:** Distributes batches balanced by `len(batch) * max_duration_in_batch`
  - **Vectorized with numpy:** ~30s to iterate 77M+ samples (vs 8min with loops)

---

### 6. CONFIGURATION FILES

**Training Config:** `/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/configs/train/stage1_prod_8xh200.yaml` (164 lines)

**Key Settings:**
```yaml
trainer:
  devices: 8
  precision: bf16-mixed
  max_steps: 200000
  val_check_interval: 10000
  accumulate_grad_batches: 4
  gradient_clip_val: 1.0

model:
  encoder:
    n_layers: 42
    d_model: 1024
    subsampling: dw_striding
    subsampling_factor: 8
    conv_kernel_size: 9

  decoder:
    pred_hidden: 640
    pred_rnn_layers: 2

  joint:
    num_extra_outputs: 5  # TDT duration heads
    fuse_loss_wer: true
    fused_batch_size: 2
    joint_hidden: 640

  loss:
    loss_name: tdt
    tdt_kwargs:
      durations: [0, 1, 2, 3, 4]
      sigma: 0.02
      omega: 0.1

  aux_ctc:
    ctc_loss_weight: 0.3  # warmed up from 0→0.3

  tokenizer:
    dir: tokenizers/stage1_prod_bpe
    type: bpe

  optim:
    name: adamw
    betas: [0.9, 0.98]
    weight_decay: 0.0001
    sched:
      name: CosineAnnealing
      warmup_steps: 2000
      min_lr: 0.00001
```

**Data Configs:**
- `/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/configs/data/stage1_prod_input_cfg.yaml` (3.5K)
- `/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/configs/data/stage1_prod_input_cfg_v2.yaml` (1.3M, per-language dataset configs)
- `/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/configs/data/local_smoke.yaml` (smoke test setup)

---

### 7. DATA PIPELINE SCRIPTS

**Key Scripts:**

1. **`scripts/build_manifest.py`** — Builds NeMo JSONL manifests from shard metadata
   - Reads metadata.parquet + audio_index.parquet
   - Outputs: `audio_filepath, text, duration, lang, taskname, source_lang, target_lang`
   - Supports hive-partitioned and flat directory layouts

2. **`scripts/split_manifest.py`** — Stratified train/val splits

3. **`scripts/build_tokenizer.py`** — Trains 32K SentencePiece BPE tokenizer
   - Quality-filtered text from all 12 languages
   - Located at: `tokenizers/stage1_prod_bpe/`
   - tokenizer.model: 962KB, vocab.txt: 472KB

4. **`scripts/generate_input_cfg.py`** — Generates NeMo input_cfg YAML

5. **Data Reading:** `lf_asr/data/tar_offset_reader.py`
   - Zero-extraction tar loading via `os.pread()` with byte offsets
   - Thread-safe fd cache (max 256 open tar files)
   - Returns (waveform, 16kHz sample rate)

6. **Dataset:** `lf_asr/data/tar_offset_dataset.py` + custom `NeMoTarOffsetDataset`
   - Wraps TarOffsetReader for NeMo training
   - Returns: (signal, signal_len, tokens, token_len, lang_id)
   - Pre-computes language IDs in constructor

---

### 8. TOKENIZER

**File Location:** `/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/tokenizers/stage1_prod_bpe/`

**Details:**
- Type: BPE (SentencePiece)
- Vocabulary Size: 32K (32,769 including special tokens)
- Coverage: 11 Indic languages + English
- Quality filtering applied during training
- Model: `tokenizer.model` (962KB)
- Vocabulary: `vocab.txt` (472KB)

**Usage:**
```python
from sentencepiece import SentencePieceProcessor
sp = SentencePieceProcessor()
sp.Load('tokenizers/stage1_prod_bpe/tokenizer.model')
token_ids = sp.EncodeAsIds(text)
```

---

### 9. HYBRID TDT ARCHITECTURE MODIFICATIONS

**TDT Loss Configuration** (`configs/train/stage1_prod_8xh200.yaml`, lines 85-91):
```yaml
loss:
  loss_name: tdt
  tdt_kwargs:
    durations: [0, 1, 2, 3, 4]
    sigma: 0.02
    omega: 0.1
    fastemit_lambda: 0.0
    clamp: -1.0
```

**Joint Network with Duration Heads** (line 74):
```yaml
joint:
  num_extra_outputs: 5  # Predicts duration 0, 1, 2, 3, or 4
```

**Why TDT over RNNT?**
- Matches Parakeet TDT 1.1B pretraining objective
- Allows skipping multiple frames per emission (more efficient)
- Better speed/accuracy tradeoff for multilingual data

---

### 10. PRETRAINED ENCODER & CHECKPOINT HANDLING

**Pretrained Encoder:**
- Parakeet TDT 1.1B (42-layer FastConformer, English pretrained)
- Loaded via `--pretrained-encoder` argument in train_prod.py
- Expected path: `pretrained/parakeet_tdt_1.1b_encoder.pt` (not present locally)

**Checkpoint Scripts:**
- `scripts/restore_checkpoint.py` — Restores from Cloudflare R2
- `scripts/upload_checkpoint.py` — Uploads to R2
- `scripts/auto_upload_checkpoints.py` — Automatic checkpoint watcher

**Checkpoint Format:**
- `.ckpt` (PyTorch Lightning checkpoint with optimizer state)
- `.nemo` (NeMo native format with config + weights)
- Saved every 5K steps with `save_top_k: 3`

---

### 11. PLANNED STAGE 2: GEMMA INTEGRATION

**From README.md & PLAN.md:**

**Stage 2: Gemma LLM Decoder (planned)** — Replaces RNN decoder with Gemma 3 1B for improved language modeling

**Architecture (Stage 2):**
```
FastConformer Encoder (1024-dim)
    ↓
Projection Layer (Conv1D stride=2 → Linear 1024→1536)
    ↓
Gemma 3 1B Decoder (1536-dim hidden, 26 layers)
```

**Gemma Details:**
- Model: `google/gemma-3-1b-pt` (1B params, 1536-dim hidden, 26 layers)
- Tokenizer: 262K SentencePiece with 4K+ Indic tokens
- Task: Transcription only (no translation)
- Language tokens: `<|hi|>`, `<|bn|>`, ..., `<|en|>`
- Task token: `<|transcribe|>`

**Training Plan:**
- Stage 2a: Train projection only (encoder + LLM frozen) — 20K steps
- Stage 2b: Full joint training (projection + LLM, encoder frozen or LoRA) — 100K-200K steps
- Stage 3: Quality fine-tuning on high-quality subset — 10K-20K steps

---

### 12. KEY FILES SUMMARY

**Core Training:**
- `scripts/train_prod.py` (1,023 lines) — Main training with TDT, DiffLR, language embedding
- `configs/train/stage1_prod_8xh200.yaml` (164 lines) — 8×H200 training config
- `lf_asr/data/production_sampler.py` — Duration bucketing + language rebalancing
- `lf_asr/data/tar_offset_reader.py` — Zero-extraction tar I/O

**Data Pipeline:**
- `scripts/build_manifest.py` — Build manifests from shards
- `scripts/split_manifest.py` — Stratified splits
- `scripts/build_tokenizer.py` — Train BPE tokenizer
- `tools/phase3_*` — Phase 3 processing (tar offset building, splits, validation)

**Configuration & Constants:**
- `src/maya_asr/config.py` — Central config (languages, paths, dims)
- `configs/data/stage1_prod_input_cfg*.yaml` — Dataset configs
- `configs/data/local_smoke.yaml` — Smoke test config

**Documentation:**
- `README.md` (10.8K) — Architecture, training, inference overview
- `PLAN.md` (22.3K) — Detailed Stage 2 planning (Gemma integration)
- `docs/phase2_runbook.md` — Data processing runbook

**Tokenizer:**
- `tokenizers/stage1_prod_bpe/tokenizer.model` — SentencePiece model
- `tokenizers/stage1_prod_bpe/vocab.txt` — 32K vocabulary

---

### 13. HARDWARE REQUIREMENTS & OPTIMIZATION

**Production Training (8×H200):**
- GPUs: 8× NVIDIA H200 (140GB HBM3)
- Precision: BF16 mixed
- Max steps: 200K (~5-7 days)
- Effective batch size: 8 GPUs × micro-batch × 4 grad accum

**VRAM Budget (per H200):**
- Model weights (bf16): ~4.5GB
- Optimizer states (AdamW fp32): ~13.5GB
- Activations + joint tensor: ~75GB budget (75% safety margin)
- Joint tensor dominates: `B × T_enc × T_dec × (V + 1 + durations)`

**OOM Protection (Triple-Cap):**
1. Per-bucket BS caps (24→6 samples)
2. max_batch_duration (120s audio)
3. max_tokens_in_batch (400 tokens) — **CRITICAL** for controlling T_dec

---

### 14. SUPPORTED LANGUAGES

12 total: Hindi (hi), Bengali (bn), Tamil (ta), Telugu (te), Marathi (mr), Gujarati (gu), Kannada (kn), Malayalam (ml), Punjabi (pa), Odia (or), Assamese (as), English (en)

---

### 15. KEY ARGUMENTS & TRAINING COMMAND

```bash
python3 scripts/train_prod.py \
  --config configs/train/stage1_prod_8xh200.yaml \
  --train-parquet artifacts/phase3/production_train_final.parquet \
  --val-manifest data/manifests/stage1_prod_val_v2.jsonl \
  --pretrained-encoder pretrained/parakeet_tdt_1.1b_encoder.pt \
  --encoder-lr 2e-5 \
  --head-lr 1e-3 \
  --freeze-encoder-steps 5000 \
  --ctc-warmup-steps 3000 \
  --max-batch-dur 120.0 \
  --max-tokens-in-batch 400 \
  --temperature 0.3
```

**Key Arguments:**
- `--encoder-lr`: 2e-5 (preserve pretrained features)
- `--head-lr`: 1e-3 (fast adaptation for new vocab)
- `--freeze-encoder-steps`: 5K (warm up heads before encoder adaptation)
- `--ctc-warmup-steps`: 3K (ramp aux loss gradually)
- `--temperature`: 0.3 (aggressive language rebalancing)
- `--no-lang-embed`: Disable language conditioning (default: enabled)