1	#!/usr/bin/env python3
2	"""
3	Smoke test: loads from config.yaml, constructs Trainer through setup_model/setup_data,
4	runs one forward/backward step, verifies decoder alignment via the real helper,
5	and runs one generate call.
6
7	Usage:
8	    python smoke_test.py --config config.yaml
9	"""
10
11	import argparse
12	import os
13	import sys
14
15	import torch
16	import numpy as np
17
18	os.environ["HF_HOME"] = "/workspace/.hf_home"
19
20	# ── Imports from the actual training code ───────────────────────────────
21	sys.path.insert(0, os.path.dirname(__file__))
22	from train import TrainConfig, Trainer, load_config
23	from tokenizer_utils import decode_tokens
24
25
26	def test_trainer_setup(config):
27	    """Construct Trainer through setup_model() — catches checkpointing, loading, etc."""
28	    print("=" * 60)
29	    print("TEST 1: Trainer.setup_model()")
30	    print("=" * 60)
31
32	    # Ensure no distributed env vars so Trainer runs in single-GPU mode
33	    for var in ["RANK", "LOCAL_RANK", "WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT"]:
34	        os.environ.pop(var, None)
35
36	    trainer = Trainer(config)
37	    trainer.setup_model()
38	    print(f"  Model loaded: {type(trainer.raw_model).__name__}")
39	    print(f"  Vocab size: {trainer.raw_model.config.vocab_size}")
40	    print(f"  build_prompt available: {hasattr(trainer.raw_model, 'build_prompt')}")
41	    return trainer
42
43
44	def test_decoder_alignment(trainer):
45	    """Verify decoder_input_ids/labels alignment using the REAL _build_decoder_inputs."""
46	    print("\n" + "=" * 60)
47	    print("TEST 2: Decoder alignment (via real _build_decoder_inputs)")
48	    print("=" * 60)
49
50	    # Set up the prompt IDs (normally done in setup_data)
51	    tokenizer = trainer.processor.tokenizer
52	    trainer.tokenizer = tokenizer
53	    trainer.lang_prompt_ids = {}
54	    for lang in ["en", "hi", "te", "ta", "ml", "bn", "gu", "kn", "pa", "mr", "or", "as"]:
55	        prompt_str = trainer.raw_model.build_prompt(language=lang, punctuation=True)
56	        trainer.lang_prompt_ids[lang] = tokenizer.encode(prompt_str, add_special_tokens=False)
57
58	    eos_id = tokenizer.convert_tokens_to_ids("<|endoftext|>")
59	    pad_id = tokenizer.convert_tokens_to_ids("<pad>")
60	    prompt_len = len(trainer.lang_prompt_ids["hi"])
61
62	    # Build a fake batch
63	    transcript = "नमस्ते दुनिया"
64	    transcript_ids = tokenizer.encode(transcript, add_special_tokens=False)
65	    n_tok = len(transcript_ids)
66
67	    max_tokens = max(n_tok, 20)
68	    labels_tensor = torch.full((2, max_tokens), -100, dtype=torch.long)
69	    labels_tensor[0, :n_tok] = torch.tensor(transcript_ids)
70	    labels_tensor[1, :n_tok] = torch.tensor(transcript_ids)
71
72	    batch = {
73	        'labels': labels_tensor,
74	        'token_lengths': torch.tensor([n_tok, n_tok]),
75	        'languages': ['hi', 'hi'],
76	    }
77
78	    # Call the REAL method
79	    decoder_input_ids, labels, decoder_attn_mask = trainer._build_decoder_inputs(batch)
80
81	    did = decoder_input_ids[0].tolist()
82	    lab = labels[0].tolist()
83
84	    # Verify the key invariant: for supervised positions, labels[i] == did[i+1]
85	    errors = []
86	    print(f"  prompt_len={prompt_len}, n_tok={n_tok}, eos={eos_id}, pad={pad_id}")
87	    print(f"  Checking labels[{prompt_len-1}..{prompt_len+n_tok-1}]:")
88
89	    for i in range(prompt_len - 1, prompt_len + n_tok):
90	        next_input = did[i + 1] if i + 1 < len(did) else -1
91	        if lab[i] == -100:
92	            continue
93	        if lab[i] == eos_id:
94	            # EOS prediction — next input is pad, that's expected
95	            print(f"    pos {i}: label=EOS({eos_id}) — OK (final prediction)")
96	            continue
97	        if lab[i] != next_input:
98	            errors.append(f"pos {i}: labels={lab[i]} != input[{i+1}]={next_input}")
99	            print(f"    pos {i}: label={lab[i]} != input[{i+1}]={next_input} — MISMATCH")
100	        else:
101	            print(f"    pos {i}: label={lab[i]} == input[{i+1}]={next_input} — OK")
102
103	    if errors:
104	        print(f"\n  FAILED — {len(errors)} mismatches")
105	        return False
106	    print("\n  PASSED — alignment correct")
107	    return True
108
109
110	def test_forward_backward(trainer):
111	    """One forward/backward step with synthetic data on the real Trainer."""
112	    print("\n" + "=" * 60)
113	    print("TEST 3: Forward/backward pass")
114	    print("=" * 60)
115
116	    if not torch.cuda.is_available():
117	        print("  SKIPPED — no CUDA")
118	        return True
119
120	    tokenizer = trainer.tokenizer
121	    eos_id = tokenizer.convert_tokens_to_ids("<|endoftext|>")
122	    prompt_ids = trainer.lang_prompt_ids["hi"]
123
124	    # Synthetic batch
125	    B, T = 2, 1000
126	    transcript_ids = [100, 200, 300, 400, 500]
127	    n_tok = len(transcript_ids)
128	    max_tokens = 20
129	    labels_tensor = torch.full((B, max_tokens), -100, dtype=torch.long)
130	    for i in range(B):
131	        labels_tensor[i, :n_tok] = torch.tensor(transcript_ids)
132
133	    batch = {
134	        'mel': torch.randn(B, 128, T, dtype=torch.float16),
135	        'mel_lengths': torch.tensor([T, T // 2]),
136	        'labels': labels_tensor,
137	        'token_lengths': torch.tensor([n_tok, n_tok]),
138	        'languages': ['hi', 'hi'],
139	    }
140
141	    # Run the real train_step
142	    trainer.model.train()
143	    trainer.optimizer = torch.optim.AdamW(trainer.raw_model.parameters(), lr=1e-5)
144	    trainer.optimizer.zero_grad()
145	    trainer.spec_augment = None
146
147	    metrics = trainer.train_step(batch)
148
149	    print(f"  loss={metrics['loss']:.4f}, batch_size={metrics['batch_size']}")
150	    has_nan = np.isnan(metrics['loss']) or np.isinf(metrics['loss'])
151	    if has_nan:
152	        print("  FAILED — NaN/Inf loss")
153	        return False
154	    print("  PASSED")
155	    return True
156
157
158	def test_generate(trainer):
159	    """One generate call on synthetic mel."""
160	    print("\n" + "=" * 60)
161	    print("TEST 4: Generate call")
162	    print("=" * 60)
163
164	    if not torch.cuda.is_available():
165	        print("  SKIPPED — no CUDA")
166	        return True
167
168	    model = trainer.raw_model
169	    tokenizer = trainer.tokenizer
170	    model.eval()
171
172	    mel = torch.randn(1, 128, 500, dtype=torch.bfloat16).to(trainer.device)
173	    mel_length = torch.tensor([500], dtype=torch.long).to(trainer.device)
174
175	    prompt_ids = trainer.lang_prompt_ids["hi"]
176	    decoder_input_ids = torch.tensor([prompt_ids], dtype=torch.long).to(trainer.device)
177
178	    with torch.no_grad():
179	        outputs = model.generate(
180	            input_features=mel,
181	            length=mel_length,
182	            decoder_input_ids=decoder_input_ids,
183	            max_new_tokens=20,
184	        )
185
186	    text = decode_tokens(tokenizer, outputs[0])
187	    print(f"  Generated {outputs.shape[1]} tokens: {text[:100]}...")
188	    print("  PASSED")
189	    return True
190
191
192	def main():
193	    parser = argparse.ArgumentParser()
194	    parser.add_argument("--config", default="config.yaml")
195	    args = parser.parse_args()
196
197	    config = load_config(args.config)
198	    print(f"Config loaded: model_name={config.model_name}")
199
200	    results = {}
201
202	    trainer = test_trainer_setup(config)
203	    results["trainer_setup"] = trainer is not None
204
205	    results["alignment"] = test_decoder_alignment(trainer)
206	    results["forward_backward"] = test_forward_backward(trainer)
207	    results["generate"] = test_generate(trainer)
208
209	    print("\n" + "=" * 60)
210	    print("SMOKE TEST SUMMARY")
211	    print("=" * 60)
212	    for name, passed in results.items():
213	        print(f"  {name}: {'PASS' if passed else 'FAIL'}")
214
215	    all_pass = all(results.values())
216	    print(f"\nOverall: {'ALL PASS' if all_pass else 'SOME FAILURES'}")
217	    return 0 if all_pass else 1
218
219
220	if __name__ == "__main__":
221	    sys.exit(main())
222

#!/usr/bin/env python3
"""
Smoke test: loads from config.yaml, constructs Trainer through setup_model/setup_data,
runs one forward/backward step, verifies decoder alignment via the real helper,
and runs one generate call.

Usage:
    python smoke_test.py --config config.yaml
"""

import argparse
import os
import sys

import torch
import numpy as np

os.environ["HF_HOME"] = "/workspace/.hf_home"

# ── Imports from the actual training code ───────────────────────────────
sys.path.insert(0, os.path.dirname(__file__))
from train import TrainConfig, Trainer, load_config
from tokenizer_utils import decode_tokens

def test_trainer_setup(config):
    """Construct Trainer through setup_model() — catches checkpointing, loading, etc."""
    print("=" * 60)
    print("TEST 1: Trainer.setup_model()")
    print("=" * 60)

    # Ensure no distributed env vars so Trainer runs in single-GPU mode
    for var in ["RANK", "LOCAL_RANK", "WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT"]:
        os.environ.pop(var, None)

    trainer = Trainer(config)
    trainer.setup_model()
    print(f"  Model loaded: {type(trainer.raw_model).__name__}")
    print(f"  Vocab size: {trainer.raw_model.config.vocab_size}")
    print(f"  build_prompt available: {hasattr(trainer.raw_model, 'build_prompt')}")
    return trainer

def test_decoder_alignment(trainer):
    """Verify decoder_input_ids/labels alignment using the REAL _build_decoder_inputs."""
    print("\n" + "=" * 60)
    print("TEST 2: Decoder alignment (via real _build_decoder_inputs)")
    print("=" * 60)

    # Set up the prompt IDs (normally done in setup_data)
    tokenizer = trainer.processor.tokenizer
    trainer.tokenizer = tokenizer
    trainer.lang_prompt_ids = {}
    for lang in ["en", "hi", "te", "ta", "ml", "bn", "gu", "kn", "pa", "mr", "or", "as"]:
        prompt_str = trainer.raw_model.build_prompt(language=lang, punctuation=True)
        trainer.lang_prompt_ids[lang] = tokenizer.encode(prompt_str, add_special_tokens=False)

    eos_id = tokenizer.convert_tokens_to_ids("<|endoftext|>")
    pad_id = tokenizer.convert_tokens_to_ids("<pad>")
    prompt_len = len(trainer.lang_prompt_ids["hi"])

    # Build a fake batch
    transcript = "नमस्ते दुनिया"
    transcript_ids = tokenizer.encode(transcript, add_special_tokens=False)
    n_tok = len(transcript_ids)

    max_tokens = max(n_tok, 20)
    labels_tensor = torch.full((2, max_tokens), -100, dtype=torch.long)
    labels_tensor[0, :n_tok] = torch.tensor(transcript_ids)
    labels_tensor[1, :n_tok] = torch.tensor(transcript_ids)

    batch = {
        'labels': labels_tensor,
        'token_lengths': torch.tensor([n_tok, n_tok]),
        'languages': ['hi', 'hi'],
    }

    # Call the REAL method
    decoder_input_ids, labels, decoder_attn_mask = trainer._build_decoder_inputs(batch)

    did = decoder_input_ids[0].tolist()
    lab = labels[0].tolist()

    # Verify the key invariant: for supervised positions, labels[i] == did[i+1]
    errors = []
    print(f"  prompt_len={prompt_len}, n_tok={n_tok}, eos={eos_id}, pad={pad_id}")
    print(f"  Checking labels[{prompt_len-1}..{prompt_len+n_tok-1}]:")

    for i in range(prompt_len - 1, prompt_len + n_tok):
        next_input = did[i + 1] if i + 1 < len(did) else -1
        if lab[i] == -100:
            continue
        if lab[i] == eos_id:
            # EOS prediction — next input is pad, that's expected
            print(f"    pos {i}: label=EOS({eos_id}) — OK (final prediction)")
            continue
        if lab[i] != next_input:
            errors.append(f"pos {i}: labels={lab[i]} != input[{i+1}]={next_input}")
            print(f"    pos {i}: label={lab[i]} != input[{i+1}]={next_input} — MISMATCH")
        else:
            print(f"    pos {i}: label={lab[i]} == input[{i+1}]={next_input} — OK")

    if errors:
        print(f"\n  FAILED — {len(errors)} mismatches")
        return False
    print("\n  PASSED — alignment correct")
    return True

def test_forward_backward(trainer):
    """One forward/backward step with synthetic data on the real Trainer."""
    print("\n" + "=" * 60)
    print("TEST 3: Forward/backward pass")
    print("=" * 60)

    if not torch.cuda.is_available():
        print("  SKIPPED — no CUDA")
        return True

    tokenizer = trainer.tokenizer
    eos_id = tokenizer.convert_tokens_to_ids("<|endoftext|>")
    prompt_ids = trainer.lang_prompt_ids["hi"]

    # Synthetic batch
    B, T = 2, 1000
    transcript_ids = [100, 200, 300, 400, 500]
    n_tok = len(transcript_ids)
    max_tokens = 20
    labels_tensor = torch.full((B, max_tokens), -100, dtype=torch.long)
    for i in range(B):
        labels_tensor[i, :n_tok] = torch.tensor(transcript_ids)

    batch = {
        'mel': torch.randn(B, 128, T, dtype=torch.float16),
        'mel_lengths': torch.tensor([T, T // 2]),
        'labels': labels_tensor,
        'token_lengths': torch.tensor([n_tok, n_tok]),
        'languages': ['hi', 'hi'],
    }

    # Run the real train_step
    trainer.model.train()
    trainer.optimizer = torch.optim.AdamW(trainer.raw_model.parameters(), lr=1e-5)
    trainer.optimizer.zero_grad()
    trainer.spec_augment = None

    metrics = trainer.train_step(batch)

    print(f"  loss={metrics['loss']:.4f}, batch_size={metrics['batch_size']}")
    has_nan = np.isnan(metrics['loss']) or np.isinf(metrics['loss'])
    if has_nan:
        print("  FAILED — NaN/Inf loss")
        return False
    print("  PASSED")
    return True

def test_generate(trainer):
    """One generate call on synthetic mel."""
    print("\n" + "=" * 60)
    print("TEST 4: Generate call")
    print("=" * 60)

    if not torch.cuda.is_available():
        print("  SKIPPED — no CUDA")
        return True

    model = trainer.raw_model
    tokenizer = trainer.tokenizer
    model.eval()

    mel = torch.randn(1, 128, 500, dtype=torch.bfloat16).to(trainer.device)
    mel_length = torch.tensor([500], dtype=torch.long).to(trainer.device)

    prompt_ids = trainer.lang_prompt_ids["hi"]
    decoder_input_ids = torch.tensor([prompt_ids], dtype=torch.long).to(trainer.device)

    with torch.no_grad():
        outputs = model.generate(
            input_features=mel,
            length=mel_length,
            decoder_input_ids=decoder_input_ids,
            max_new_tokens=20,
        )

    text = decode_tokens(tokenizer, outputs[0])
    print(f"  Generated {outputs.shape[1]} tokens: {text[:100]}...")
    print("  PASSED")
    return True

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default="config.yaml")
    args = parser.parse_args()

    config = load_config(args.config)
    print(f"Config loaded: model_name={config.model_name}")

    results = {}

    trainer = test_trainer_setup(config)
    results["trainer_setup"] = trainer is not None

    results["alignment"] = test_decoder_alignment(trainer)
    results["forward_backward"] = test_forward_backward(trainer)
    results["generate"] = test_generate(trainer)

    print("\n" + "=" * 60)
    print("SMOKE TEST SUMMARY")
    print("=" * 60)
    for name, passed in results.items():
        print(f"  {name}: {'PASS' if passed else 'FAIL'}")

    all_pass = all(results.values())
    print(f"\nOverall: {'ALL PASS' if all_pass else 'SOME FAILURES'}")
    return 0 if all_pass else 1

if __name__ == "__main__":
    sys.exit(main())


Now I understand the model's inference pattern. Let me check what's needed for the processor to handle raw audio.


Bash

ls /home/ubuntu/training/cohere-transcribe-indic/tokenizer_extension/extended_model/ | head -20


config.json
configuration_cohere_asr.py
generation_config.json
modeling_cohere_asr.py
preprocessor_config.json
processing_cohere_asr.py
processor_config.json
special_tokens_map.json
tokenization_cohere_asr.py
tokenizer_config.json
tokenizer.json
vocab_info.json


1	# Cohere Transcribe — Indic Language Extension
2
3	Finetuning [CohereLabs/cohere-transcribe-03-2026](https://huggingface.co/CohereLabs/cohere-transcribe-03-2026) (2B Fast-Conformer encoder + Transformer decoder) for 11 Indic languages.
4
5	## Overview
6
7	| Metric | Value |
8	|--------|-------|
9	| Base model | CohereLabs/cohere-transcribe-03-2026 (2.15B params) |
10	| Languages | English + Hindi, Telugu, Tamil, Malayalam, Bengali, Gujarati, Kannada, Punjabi, Marathi, Odia, Assamese |
11	| Training data | 73M utterances, 145K hours (cleaned from 74.5M / 147K hours) |
12	| Tokenizer | Extended from 16K → 59,554 vocab (32K Indic subwords added) |
13	| Hardware | 8x NVIDIA H200 (141GB HBM3e each) |
14	| Training time | ~5 days (2 epochs) |
15
16	## Tokenizer Extension
17
18	The original 16K BPE tokenizer had zero Indic subwords — every Indic character fell back to 3 UTF-8 bytes (~13 tokens/word for Hindi vs ~1.2 for English).
19
20	We trained a 32K SentencePiece unigram model on 57M Indic transcripts, then merged the new subwords into the existing BPE tokenizer via synthetic merge rules.
21
22	**Result: 89.5% token reduction for Indic text** (13 tok/word → 1.5 tok/word), with zero impact on English tokenization.
23
24	| Language | Old tok/word | New tok/word | Reduction |
25	|----------|-------------|--------------|-----------|
26	| Hindi | 11.95 | 1.22 | 89.8% |
27	| Telugu | 16.58 | 1.54 | 90.7% |
28	| Tamil | 22.00 | 1.91 | 91.3% |
29	| Bengali | 13.97 | 1.18 | 91.6% |
30	| English | 1.17 | 1.17 | 0.0% |
31
32	## Data Pipeline
33
34	### Data Cleaning (applied to training manifest)
35	- Removed empty/short transcripts, REDACTED, `<unintelligible>` tags
36	- Stripped `[singing]`, `[laugh]`, `[noise]` annotations
37	- Removed segments < 0.5s or > 30s
38	- Removed very slow speech rate (< 2 chars/sec, likely misaligned)
39	- Normalized ellipses, zero-width spaces, native digits → Arabic
40	- Fixed language tag mismatches (relabeled Telugu-in-wrong-lang, dropped ambiguous Devanagari)
41	- Deduplicated cross-language and within-shard duplicates
42	- **Net: 72,998,025 rows / 145,139 hours (2.06% removed)**
43
44	### Pre-computed Mel Spectrograms
45	- 128-channel log-mel, 16kHz, 100fps, per-feature normalized
46	- Extracted from tar shards to individual `.npy` files for fast random access
47	- Per-shard parquet indexes for streaming reads
48
49	## Training Configuration
50
51	```yaml
52	model: 59,554 vocab (16K original + 43K Indic subwords + intermediates)
53	batch: 128 utterances/GPU × 8 GPUs = 1024 effective batch
54	lr: 2e-4, cosine decay, 2% warmup
55	epochs: 2
56	precision: bf16
57	strategy: DDP
58	gradient_checkpointing: manual (encoder conformer layers)
59	label_smoothing: 0.1 (via NLLLoss — model head outputs log-probs)
60	```
61
62	### Key Implementation Details
63	- **Decoder prompt**: 9 atomic control tokens per sample (language-conditioned)
64	- **Label shift**: Model does NOT shift labels internally — `labels[i] = decoder_input_ids[i+1]`
65	- **Loss**: NLLLoss (not CrossEntropyLoss) because model head applies log_softmax
66	- **Weight tying**: Decoder embedding and output head share weights
67	- **Temperature sampling**: τ=5.0 for language balancing (upsamples Assamese/Odia)
68
69	## Repository Structure
70
71	```
72	training/
73	├── train.py                 # Main trainer (DDP, grad checkpointing, wandb)
74	├── dataset.py               # Streaming dataset (tar + extracted .npy support)
75	├── dataset_fast.py           # Fast dataset (direct .npy reads, bucketed batching)
76	├── config.yaml              # Production training config
77	├── config_smoke.yaml        # Smoke test config (small subset)
78	├── evaluate.py              # WER/CER evaluation script
79	├── tokenizer_utils.py       # Extended tokenizer loader + decode helper
80	├── smoke_test.py            # Startup validation (alignment, forward/backward, generate)
81	├── resize_embeddings.py     # Embedding resize + weight tying for extended vocab
82	├── extend_tokenizer.py      # SentencePiece training + tokenizer extension pipeline
83	├── merge_tokenizer_final.py # BPE merge with synthetic merge rules
84	├── push_to_r2.sh            # Checkpoint upload to Cloudflare R2
85	└── tokenizer_extension/
86	    └── extended_model/      # Extended checkpoint (59K vocab, bf16)
87	        ├── model.safetensors
88	        ├── config.json
89	        ├── tokenizer.json
90	        └── modeling_cohere_asr.py
91
92	maya-asr/
93	├── manifests/
94	│   ├── training_manifest_cleaned.parquet  # 73M rows, cleaned
95	│   └── training_manifest_smoke.parquet    # 60K rows, smoke test
96	├── mel_extracted/           # Extracted mel .npy files (from tar shards)
97	├── mel_extracted_index/     # Per-shard parquet indexes
98	├── clean_manifest.py        # Data cleanup pipeline
99	└── build_shard_indexes.py   # Per-shard index builder
100	```

# Cohere Transcribe — Indic Language Extension

Finetuning [CohereLabs/cohere-transcribe-03-2026](https://huggingface.co/CohereLabs/cohere-transcribe-03-2026) (2B Fast-Conformer encoder + Transformer decoder) for 11 Indic languages.

## Overview

| Metric | Value |
|--------|-------|
| Base model | CohereLabs/cohere-transcribe-03-2026 (2.15B params) |
| Languages | English + Hindi, Telugu, Tamil, Malayalam, Bengali, Gujarati, Kannada, Punjabi, Marathi, Odia, Assamese |
| Training data | 73M utterances, 145K hours (cleaned from 74.5M / 147K hours) |
| Tokenizer | Extended from 16K → 59,554 vocab (32K Indic subwords added) |
| Hardware | 8x NVIDIA H200 (141GB HBM3e each) |
| Training time | ~5 days (2 epochs) |

## Tokenizer Extension

The original 16K BPE tokenizer had zero Indic subwords — every Indic character fell back to 3 UTF-8 bytes (~13 tokens/word for Hindi vs ~1.2 for English).

We trained a 32K SentencePiece unigram model on 57M Indic transcripts, then merged the new subwords into the existing BPE tokenizer via synthetic merge rules.

**Result: 89.5% token reduction for Indic text** (13 tok/word → 1.5 tok/word), with zero impact on English tokenization.

| Language | Old tok/word | New tok/word | Reduction |
|----------|-------------|--------------|-----------|
| Hindi | 11.95 | 1.22 | 89.8% |
| Telugu | 16.58 | 1.54 | 90.7% |
| Tamil | 22.00 | 1.91 | 91.3% |
| Bengali | 13.97 | 1.18 | 91.6% |
| English | 1.17 | 1.17 | 0.0% |

## Data Pipeline

### Data Cleaning (applied to training manifest)
- Removed empty/short transcripts, REDACTED, `<unintelligible>` tags
- Stripped `[singing]`, `[laugh]`, `[noise]` annotations
- Removed segments < 0.5s or > 30s
- Removed very slow speech rate (< 2 chars/sec, likely misaligned)
- Normalized ellipses, zero-width spaces, native digits → Arabic
- Fixed language tag mismatches (relabeled Telugu-in-wrong-lang, dropped ambiguous Devanagari)
- Deduplicated cross-language and within-shard duplicates
- **Net: 72,998,025 rows / 145,139 hours (2.06% removed)**

### Pre-computed Mel Spectrograms
- 128-channel log-mel, 16kHz, 100fps, per-feature normalized
- Extracted from tar shards to individual `.npy` files for fast random access
- Per-shard parquet indexes for streaming reads

## Training Configuration

```yaml
model: 59,554 vocab (16K original + 43K Indic subwords + intermediates)
batch: 128 utterances/GPU × 8 GPUs = 1024 effective batch
lr: 2e-4, cosine decay, 2% warmup
epochs: 2
precision: bf16
strategy: DDP
gradient_checkpointing: manual (encoder conformer layers)
label_smoothing: 0.1 (via NLLLoss — model head outputs log-probs)
```

### Key Implementation Details
- **Decoder prompt**: 9 atomic control tokens per sample (language-conditioned)
- **Label shift**: Model does NOT shift labels internally — `labels[i] = decoder_input_ids[i+1]`
- **Loss**: NLLLoss (not CrossEntropyLoss) because model head applies log_softmax
- **Weight tying**: Decoder embedding and output head share weights
- **Temperature sampling**: τ=5.0 for language balancing (upsamples Assamese/Odia)

## Repository Structure

```
training/
├── train.py                 # Main trainer (DDP, grad checkpointing, wandb)
├── dataset.py               # Streaming dataset (tar + extracted .npy support)
├── dataset_fast.py           # Fast dataset (direct .npy reads, bucketed batching)
├── config.yaml              # Production training config
├── config_smoke.yaml        # Smoke test config (small subset)
├── evaluate.py              # WER/CER evaluation script
├── tokenizer_utils.py       # Extended tokenizer loader + decode helper
├── smoke_test.py            # Startup validation (alignment, forward/backward, generate)
├── resize_embeddings.py     # Embedding resize + weight tying for extended vocab
├── extend_tokenizer.py      # SentencePiece training + tokenizer extension pipeline
├── merge_tokenizer_final.py # BPE merge with synthetic merge rules
├── push_to_r2.sh            # Checkpoint upload to Cloudflare R2
└── tokenizer_extension/
    └── extended_model/      # Extended checkpoint (59K vocab, bf16)
        ├── model.safetensors
        ├── config.json
        ├── tokenizer.json
        └── modeling_cohere_asr.py

maya-asr/
├── manifests/
│   ├── training_manifest_cleaned.parquet  # 73M rows, cleaned
│   └── training_manifest_smoke.parquet    # 60K rows, smoke test
├── mel_extracted/           # Extracted mel .npy files (from tar shards)
├── mel_extracted_index/     # Per-shard parquet indexes
├── clean_manifest.py        # Data cleanup pipeline
└── build_shard_indexes.py   # Per-shard index builder
```


cat /home/ubuntu/training/cohere-transcribe-indic/tokenizer_extension/extended_model/processing_cohere_asr.py 2>/dev/null | head -50 || echo "not found"


import logging
import math
import random
from pathlib import Path

import librosa
import numpy as np
import torch
import torch.nn.functional as F
from safetensors.torch import load_file as safetensors_load_file
from torch import nn
from transformers import AutoFeatureExtractor, AutoTokenizer, BatchFeature
from transformers.feature_extraction_sequence_utils import SequenceFeatureExtractor
from transformers.processing_utils import ProcessorMixin

from .configuration_cohere_asr import _dynamo_disable

logger = logging.getLogger(__name__)

DITHER_CONSTANT = 1e-5

class FilterbankFeatures(nn.Module):
    """Filterbank features extraction module.

    Follows NeMo's FilterbankFeatures implementation.
    """

    window: torch.Tensor
    fb: torch.Tensor

    def __init__(
        self,
        sample_rate=16000,
        n_window_size=320,
        n_window_stride=160,
        window="hann",
        normalize="per_feature",
        n_fft=None,
        preemph=0.97,
        nfilt=64,
        lowfreq=0,
        highfreq=None,
        log=True,
        log_zero_guard_type="add",
        log_zero_guard_value=2**-24,
        dither=DITHER_CONSTANT,
        pad_to=16,
        max_duration=30,
        frame_splicing=1,


<task-notification>
<task-id>b24ms6xfr</task-id>
<tool-use-id>toolu_01BQdQJKiUxnCN7Bbtppkr3d</tool-use-id>
<output-file>/tmp/claude-1001/-home-ubuntu-training/ed45c711-ad1f-4a43-a0e4-94f0e1c224d8/tasks/b24ms6xfr.output</output-file>
<status>completed</status>
<summary>Background command "Download model and processor from R2 to local checkpoint dir" completed (exit code 0)</summary>
</task-notification>


Good. The model uses a custom processor that extracts 128-channel mel features internally. Now let me look at how the processor handles raw audio for inference.


grep -n "def __call__\|def process\|class CohereAsr.*Processor\|class CohereAsr.*Extractor" /home/ubuntu/training/cohere-transcribe-indic/tokenizer_extension/extended_model/processing_cohere_asr.py | head -20


303:class CohereAsrFeatureExtractor(SequenceFeatureExtractor):
396:    def __call__(
444:class CohereAsrProcessor(ProcessorMixin):
488:    def __call__(


303	class CohereAsrFeatureExtractor(SequenceFeatureExtractor):
304	    """HF-compatible feature extractor wrapping FilterbankFeatures."""
305
306	    model_input_names = ["input_features"]
307
308	    def __init__(
309	        self,
310	        feature_size=64,
311	        sampling_rate=16000,
312	        padding_value=0.0,
313	        max_duration=30,
314	        n_window_size=320,
315	        n_window_stride=160,
316	        window="hann",
317	        normalize="per_feature",
318	        n_fft=None,
319	        preemph=0.97,
320	        lowfreq=0,
321	        highfreq=None,
322	        log=True,
323	        log_zero_guard_type="add",
324	        log_zero_guard_value=2**-24,
325	        dither=DITHER_CONSTANT,
326	        pad_to=16,
327	        frame_splicing=1,
328	        exact_pad=False,
329	        mag_power=2.0,
330	        nb_augmentation_prob=0.0,
331	        nb_max_freq=4000,
332	        mel_norm="slaney",
333	        stft_exact_pad=False,
334	        stft_conv=False,
335	        device="cpu",
336	        **kwargs,
337	    ):
338	        super().__init__(
339	            feature_size=feature_size,
340	            sampling_rate=sampling_rate,
341	            padding_value=padding_value,
342	            **kwargs,
343	        )
344	        self.max_duration = max_duration
345	        self.hop_length = n_window_stride
346	        self._device = str(device)
347	        self._fb_config = dict(
348	            sample_rate=sampling_rate,
349	            n_window_size=n_window_size,
350	            n_window_stride=n_window_stride,
351	            window=window,
352	            normalize=normalize,
353	            n_fft=n_fft,
354	            preemph=preemph,
355	            nfilt=feature_size,
356	            lowfreq=lowfreq,
357	            highfreq=highfreq,
358	            log=log,
359	            log_zero_guard_type=log_zero_guard_type,
360	            log_zero_guard_value=log_zero_guard_value,
361	            dither=dither,
362	            pad_to=pad_to,
363	            max_duration=max_duration,
364	            frame_splicing=frame_splicing,
365	            exact_pad=exact_pad,
366	            pad_value=padding_value,
367	            mag_power=mag_power,
368	            nb_augmentation_prob=nb_augmentation_prob,
369	            nb_max_freq=nb_max_freq,
370	            mel_norm=mel_norm,
371	            stft_exact_pad=stft_exact_pad,
372	            stft_conv=stft_conv,
373	            device=device,
374	        )
375	        self._filterbank = None
376
377	    @classmethod
378	    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
379	        fe = super().from_pretrained(pretrained_model_name_or_path, **kwargs)
380	        model_dir = Path(pretrained_model_name_or_path)
381	        if model_dir.is_dir():
382	            _maybe_load_preprocessor_buffers_from_checkpoint(feature_extractor=fe, model_dir=model_dir)
383	        return fe
384
385	    @property
386	    def filterbank(self):
387	        if self._filterbank is None:
388	            fb = FilterbankFeatures(**self._fb_config)
389	            fb.eval()
390	            self._filterbank = fb.to(self._device)
391	        return self._filterbank
392
393	    def get_seq_len(self, seq_len):
394	        return self.filterbank.get_seq_len(seq_len)
395
396	    def __call__(
397	        self,
398	        raw_speech,
399	        sampling_rate=None,
400	        return_tensors=None,
401	        **kwargs,
402	    ):

class CohereAsrFeatureExtractor(SequenceFeatureExtractor):
    """HF-compatible feature extractor wrapping FilterbankFeatures."""

    model_input_names = ["input_features"]

    def __init__(
        self,
        feature_size=64,
        sampling_rate=16000,
        padding_value=0.0,
        max_duration=30,
        n_window_size=320,
        n_window_stride=160,
        window="hann",
        normalize="per_feature",
        n_fft=None,
        preemph=0.97,
        lowfreq=0,
        highfreq=None,
        log=True,
        log_zero_guard_type="add",
        log_zero_guard_value=2**-24,
        dither=DITHER_CONSTANT,
        pad_to=16,
        frame_splicing=1,
        exact_pad=False,
        mag_power=2.0,
        nb_augmentation_prob=0.0,
        nb_max_freq=4000,
        mel_norm="slaney",
        stft_exact_pad=False,
        stft_conv=False,
        device="cpu",
        **kwargs,
    ):
        super().__init__(
            feature_size=feature_size,
            sampling_rate=sampling_rate,
            padding_value=padding_value,
            **kwargs,
        )
        self.max_duration = max_duration
        self.hop_length = n_window_stride
        self._device = str(device)
        self._fb_config = dict(
            sample_rate=sampling_rate,
            n_window_size=n_window_size,
            n_window_stride=n_window_stride,
            window=window,
            normalize=normalize,
            n_fft=n_fft,
            preemph=preemph,
            nfilt=feature_size,
            lowfreq=lowfreq,
            highfreq=highfreq,
            log=log,
            log_zero_guard_type=log_zero_guard_type,
            log_zero_guard_value=log_zero_guard_value,
            dither=dither,
            pad_to=pad_to,
            max_duration=max_duration,
            frame_splicing=frame_splicing,
            exact_pad=exact_pad,
            pad_value=padding_value,
            mag_power=mag_power,
            nb_augmentation_prob=nb_augmentation_prob,
            nb_max_freq=nb_max_freq,
            mel_norm=mel_norm,
            stft_exact_pad=stft_exact_pad,
            stft_conv=stft_conv,
            device=device,
        )
        self._filterbank = None

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
        fe = super().from_pretrained(pretrained_model_name_or_path, **kwargs)
        model_dir = Path(pretrained_model_name_or_path)
        if model_dir.is_dir():
            _maybe_load_preprocessor_buffers_from_checkpoint(feature_extractor=fe, model_dir=model_dir)
        return fe

    @property
    def filterbank(self):
        if self._filterbank is None:
            fb = FilterbankFeatures(**self._fb_config)
            fb.eval()
            self._filterbank = fb.to(self._device)
        return self._filterbank

    def get_seq_len(self, seq_len):
        return self.filterbank.get_seq_len(seq_len)

    def __call__(
        self,
        raw_speech,
        sampling_rate=None,
        return_tensors=None,
        **kwargs,
    ):


396	    def __call__(
397	        self,
398	        raw_speech,
399	        sampling_rate=None,
400	        return_tensors=None,
401	        **kwargs,
402	    ):
403	        """Extract mel features from raw waveform input."""
404	        if sampling_rate is not None and int(sampling_rate) != int(self.sampling_rate):
405	            raise ValueError(f"Expected sampling_rate={self.sampling_rate}, got {sampling_rate}")
406
407	        if isinstance(raw_speech, np.ndarray):
408	            if raw_speech.ndim == 1:
409	                raw_speech = [raw_speech]
410	            else:
411	                raw_speech = [s for s in raw_speech]
412	        elif isinstance(raw_speech, torch.Tensor):
413	            if raw_speech.ndim == 1:
414	                raw_speech = [raw_speech.detach().cpu().numpy()]
415	            else:
416	                raw_speech = [s.detach().cpu().numpy() for s in raw_speech]
417	        elif not isinstance(raw_speech, (list, tuple)):
418	            raise TypeError("raw_speech must be an array/tensor or list of arrays.")
419
420	        normalized = []
421	        for sample in raw_speech:
422	            arr = np.asarray(sample, dtype=np.float32)
423	            if arr.ndim != 1:
424	                raise ValueError("Each audio sample must be 1D waveform.")
425	            normalized.append(arr)
426
427	        seq_len = torch.tensor([s.shape[0] for s in normalized], dtype=torch.long)
428	        max_len = max(s.shape[0] for s in normalized)
429	        padded = np.zeros((len(normalized), max_len), dtype=np.float32)
430	        for i, s in enumerate(normalized):
431	            padded[i, : s.shape[0]] = s
432
433	        audio_tensor = torch.from_numpy(padded).to(self._device)
434	        seq_len = seq_len.to(self._device)
435	        with torch.no_grad():
436	            input_features, length = self.filterbank(audio_tensor, seq_len)
437
438	        result = BatchFeature({"input_features": input_features.cpu(), "length": length.cpu()})
439	        if return_tensors is not None:
440	            result = result.convert_to_tensors(return_tensors)
441	        return result
442
443
444	class CohereAsrProcessor(ProcessorMixin):
445	    """HF-compatible processor for Cohere ASR.
446
447	    ``ProcessorMixin._get_arguments_from_pretrained`` resolves sub-component
448	    class names by looking them up inside the ``transformers`` package, which
449	    fails for custom remote-code classes.  We override ``from_pretrained`` to
450	    use ``AutoFeatureExtractor`` / ``AutoTokenizer`` instead -- those honour
451	    ``auto_map`` and ``trust_remote_code``.
452	    """
453
454	    attributes = ["feature_extractor", "tokenizer"]
455	    feature_extractor_class = "CohereAsrFeatureExtractor"
456	    tokenizer_class = "CohereAsrTokenizer"
457
458	    def __init__(self, feature_extractor=None, tokenizer=None, **kwargs):
459	        if feature_extractor is None:
460	            raise ValueError(
461	                "CohereAsrProcessor requires a CohereAsrFeatureExtractor instance. " "Got feature_extractor=None."
462	            )
463	        if tokenizer is None:
464	            raise ValueError("CohereAsrProcessor requires a CohereAsrTokenizer instance. " "Got tokenizer=None.")
465	        # Bypass super().__init__ which calls get_possibly_dynamic_module to
466	        # validate sub-component types.  That lookup searches the transformers
467	        # package namespace and fails for remote-code classes.  We set the
468	        # attributes directly instead -- the type checks above are sufficient.
469	        self.feature_extractor = feature_extractor
470	        self.tokenizer = tokenizer
471	        self.chat_template = kwargs.get("chat_template", None)
472
473	    @classmethod
474	    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
475	        trust_remote_code = kwargs.pop("trust_remote_code", True)
476	        feature_extractor = AutoFeatureExtractor.from_pretrained(
477	            pretrained_model_name_or_path,
478	            trust_remote_code=trust_remote_code,
479	            **kwargs,
480	        )
481	        tokenizer = AutoTokenizer.from_pretrained(
482	            pretrained_model_name_or_path,
483	            trust_remote_code=trust_remote_code,
484	            **kwargs,
485	        )
486	        return cls(feature_extractor=feature_extractor, tokenizer=tokenizer)
487
488	    def __call__(
489	        self,
490	        audio=None,
491	        text=None,
492	        sampling_rate=None,
493	        return_tensors=None,
494	        **kwargs,
495	    ):
496	        """Run audio feature extraction and optional text tokenization."""
497	        if audio is None:
498	            raise ValueError("audio is required for CohereAsrProcessor.")
499
500	        result = self.feature_extractor(audio, sampling_rate=sampling_rate, return_tensors=return_tensors)
501
502	        if text is not None:
503	            add_special_tokens = kwargs.pop("add_special_tokens", False)
504	            text_inputs = self.tokenizer(
505	                text,
506	                return_tensors=return_tensors,
507	                add_special_tokens=add_special_tokens,
508	                **kwargs,
509	            )
510	            result["input_ids"] = text_inputs["input_ids"]
511	            if "attention_mask" in text_inputs:
512	                result["attention_mask"] = text_inputs["attention_mask"]
513	        return result
514
515	    def batch_decode(self, *args, **kwargs):

    def __call__(
        self,
        raw_speech,
        sampling_rate=None,
        return_tensors=None,
        **kwargs,
    ):
        """Extract mel features from raw waveform input."""
        if sampling_rate is not None and int(sampling_rate) != int(self.sampling_rate):
            raise ValueError(f"Expected sampling_rate={self.sampling_rate}, got {sampling_rate}")

        if isinstance(raw_speech, np.ndarray):
            if raw_speech.ndim == 1:
                raw_speech = [raw_speech]
            else:
                raw_speech = [s for s in raw_speech]
        elif isinstance(raw_speech, torch.Tensor):
            if raw_speech.ndim == 1:
                raw_speech = [raw_speech.detach().cpu().numpy()]
            else:
                raw_speech = [s.detach().cpu().numpy() for s in raw_speech]
        elif not isinstance(raw_speech, (list, tuple)):
            raise TypeError("raw_speech must be an array/tensor or list of arrays.")

        normalized = []
        for sample in raw_speech:
            arr = np.asarray(sample, dtype=np.float32)
            if arr.ndim != 1:
                raise ValueError("Each audio sample must be 1D waveform.")
            normalized.append(arr)

        seq_len = torch.tensor([s.shape[0] for s in normalized], dtype=torch.long)
        max_len = max(s.shape[0] for s in normalized)
        padded = np.zeros((len(normalized), max_len), dtype=np.float32)
        for i, s in enumerate(normalized):
            padded[i, : s.shape[0]] = s

        audio_tensor = torch.from_numpy(padded).to(self._device)
        seq_len = seq_len.to(self._device)
        with torch.no_grad():
            input_features, length = self.filterbank(audio_tensor, seq_len)

        result = BatchFeature({"input_features": input_features.cpu(), "length": length.cpu()})
        if return_tensors is not None:
            result = result.convert_to_tensors(return_tensors)
        return result

class CohereAsrProcessor(ProcessorMixin):
    """HF-compatible processor for Cohere ASR.

    ``ProcessorMixin._get_arguments_from_pretrained`` resolves sub-component
    class names by looking them up inside the ``transformers`` package, which
    fails for custom remote-code classes.  We override ``from_pretrained`` to
    use ``AutoFeatureExtractor`` / ``AutoTokenizer`` instead -- those honour
    ``auto_map`` and ``trust_remote_code``.
    """

    attributes = ["feature_extractor", "tokenizer"]
    feature_extractor_class = "CohereAsrFeatureExtractor"
    tokenizer_class = "CohereAsrTokenizer"

    def __init__(self, feature_extractor=None, tokenizer=None, **kwargs):
        if feature_extractor is None:
            raise ValueError(
                "CohereAsrProcessor requires a CohereAsrFeatureExtractor instance. " "Got feature_extractor=None."
            )
        if tokenizer is None:
            raise ValueError("CohereAsrProcessor requires a CohereAsrTokenizer instance. " "Got tokenizer=None.")
        # Bypass super().__init__ which calls get_possibly_dynamic_module to
        # validate sub-component types.  That lookup searches the transformers
        # package namespace and fails for remote-code classes.  We set the
        # attributes directly instead -- the type checks above are sufficient.
        self.feature_extractor = feature_extractor
        self.tokenizer = tokenizer
        self.chat_template = kwargs.get("chat_template", None)

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
        trust_remote_code = kwargs.pop("trust_remote_code", True)
        feature_extractor = AutoFeatureExtractor.from_pretrained(
            pretrained_model_name_or_path,
            trust_remote_code=trust_remote_code,
            **kwargs,
        )
        tokenizer = AutoTokenizer.from_pretrained(
            pretrained_model_name_or_path,
            trust_remote_code=trust_remote_code,
            **kwargs,
        )
        return cls(feature_extractor=feature_extractor, tokenizer=tokenizer)

    def __call__(
        self,
        audio=None,
        text=None,
        sampling_rate=None,
        return_tensors=None,
        **kwargs,
    ):
        """Run audio feature extraction and optional text tokenization."""
        if audio is None:
            raise ValueError("audio is required for CohereAsrProcessor.")

        result = self.feature_extractor(audio, sampling_rate=sampling_rate, return_tensors=return_tensors)

        if text is not None:
            add_special_tokens = kwargs.pop("add_special_tokens", False)
            text_inputs = self.tokenizer(
                text,
                return_tensors=return_tensors,
                add_special_tokens=add_special_tokens,
                **kwargs,
            )
            result["input_ids"] = text_inputs["input_ids"]
            if "attention_mask" in text_inputs:
                result["attention_mask"] = text_inputs["attention_mask"]
        return result

    def batch_decode(self, *args, **kwargs):


Now I have a complete understanding. Let me also check how `build_prompt` works.


grep -n "def build_prompt\|build_prompt" /home/ubuntu/training/cohere-transcribe-indic/tokenizer_extension/extended_model/modeling_cohere_asr.py | head -10


980:    def build_prompt(self, language: str, punctuation: bool = True) -> str:
1114:        prompt_text = self.build_prompt(language=language, punctuation=punctuation)


975	        supported_languages = set(getattr(self.config, "supported_languages", []))
976	        if language not in supported_languages:
977	            supported_joined = ", ".join(sorted(supported_languages))
978	            raise ValueError(f"Unsupported language '{language}'. Supported languages: {supported_joined}.")
979
980	    def build_prompt(self, language: str, punctuation: bool = True) -> str:
981	        """Build the decoder prompt prefix for language and punctuation settings."""
982	        pnc_token = "<|pnc|>" if punctuation else "<|nopnc|>"
983	        task_token = "<|noitn|>"
984	        return (
985	            "<|startofcontext|><|startoftranscript|><|emo:undefined|>"
986	            f"<|{language}|><|{language}|>{pnc_token}{task_token}<|notimestamp|><|nodiarize|>"
987	        )
988
989	    def _load_and_resample_audio(
990	        self,
991	        target_sample_rate: int,
992	        audio_file: Optional[str] = None,
993	        audio_array: Optional[np.ndarray] = None,
994	        sample_rate: Optional[int] = None,
995	    ) -> tuple[np.ndarray, int]:
996	        if (audio_file is None) == (audio_array is None):
997	            raise ValueError("Exactly one of audio_file or audio_array must be provided.")
998
999	        if audio_file is not None:
1000	            audio, loaded_sample_rate = sf.read(audio_file)
1001	            arr = np.asarray(audio, dtype=np.float32)
1002	            sample_rate_int = int(loaded_sample_rate)
1003	        else:
1004	            if sample_rate is None:
1005	                raise ValueError("sample_rate is required when audio_array is provided.")
1006	            arr = np.asarray(audio_array, dtype=np.float32)
1007	            sample_rate_int = int(sample_rate)
1008
1009	        if arr.ndim > 1:
1010	            arr = arr.mean(axis=1)
1011	        if arr.ndim != 1:
1012	            raise ValueError(f"Expected mono waveform (1D), got shape={arr.shape}")
1013
1014	        if sample_rate_int != target_sample_rate:
1015	            arr = librosa.resample(
1016	                arr,
1017	                orig_sr=sample_rate_int,
1018	                target_sr=target_sample_rate,
1019	            ).astype(np.float32, copy=False)
1020	            sample_rate_int = target_sample_rate
1021
1022	        return arr, sample_rate_int
1023
1024	    def _prepare_segments(

        supported_languages = set(getattr(self.config, "supported_languages", []))
        if language not in supported_languages:
            supported_joined = ", ".join(sorted(supported_languages))
            raise ValueError(f"Unsupported language '{language}'. Supported languages: {supported_joined}.")

    def build_prompt(self, language: str, punctuation: bool = True) -> str:
        """Build the decoder prompt prefix for language and punctuation settings."""
        pnc_token = "<|pnc|>" if punctuation else "<|nopnc|>"
        task_token = "<|noitn|>"
        return (
            "<|startofcontext|><|startoftranscript|><|emo:undefined|>"
            f"<|{language}|><|{language}|>{pnc_token}{task_token}<|notimestamp|><|nodiarize|>"
        )

    def _load_and_resample_audio(
        self,
        target_sample_rate: int,
        audio_file: Optional[str] = None,
        audio_array: Optional[np.ndarray] = None,
        sample_rate: Optional[int] = None,
    ) -> tuple[np.ndarray, int]:
        if (audio_file is None) == (audio_array is None):
            raise ValueError("Exactly one of audio_file or audio_array must be provided.")

        if audio_file is not None:
            audio, loaded_sample_rate = sf.read(audio_file)
            arr = np.asarray(audio, dtype=np.float32)
            sample_rate_int = int(loaded_sample_rate)
        else:
            if sample_rate is None:
                raise ValueError("sample_rate is required when audio_array is provided.")
            arr = np.asarray(audio_array, dtype=np.float32)
            sample_rate_int = int(sample_rate)

        if arr.ndim > 1:
            arr = arr.mean(axis=1)
        if arr.ndim != 1:
            raise ValueError(f"Expected mono waveform (1D), got shape={arr.shape}")

        if sample_rate_int != target_sample_rate:
            arr = librosa.resample(
                arr,
                orig_sr=sample_rate_int,
                target_sr=target_sample_rate,
            ).astype(np.float32, copy=False)
            sample_rate_int = target_sample_rate

        return arr, sample_rate_int

    def _prepare_segments(


Now let me check what the `build_prompt` language format is — the `language` parameter should use language codes like "en", "hi" etc.


121:  "supported_languages": [


121	  "supported_languages": [
122	    "en", "fr", "de", "es", "it", "pt", "nl", "pl", "el", "ar", "ja", "zh", "vi", "ko",
123	    "hi", "te", "ta", "ml", "bn", "gu", "kn", "pa", "mr", "or", "as"
124	  ],
125	  "transf_decoder": {
126	    "config_dict": {
127	      "attn_layer_dropout": 0,
128	      "attn_score_dropout": 0,
129	      "embedding_dropout": 0,
130	      "ffn_dropout": 0,
131	      "hidden_act": "relu",
132	      "hidden_size": 1024,
133	      "inner_size": 4096,
134	      "learn_positional_encodings": false,
135	      "lm_dec_hidden": 1280,
136	      "max_sequence_length": 1024,
137	      "num_attention_heads": 8,
138	      "num_layers": 8,
139	      "num_token_types": 0,
140	      "pre_ln": true,

  "supported_languages": [
    "en", "fr", "de", "es", "it", "pt", "nl", "pl", "el", "ar", "ja", "zh", "vi", "ko",
    "hi", "te", "ta", "ml", "bn", "gu", "kn", "pa", "mr", "or", "as"
  ],
  "transf_decoder": {
    "config_dict": {
      "attn_layer_dropout": 0,
      "attn_score_dropout": 0,
      "embedding_dropout": 0,
      "ffn_dropout": 0,
      "hidden_act": "relu",
      "hidden_size": 1024,
      "inner_size": 4096,
      "learn_positional_encodings": false,
      "lm_dec_hidden": 1280,
      "max_sequence_length": 1024,
      "num_attention_heads": 8,
      "num_layers": 8,
      "num_token_types": 0,
      "pre_ln": true,


Now I have everything I need. The model uses language codes like "en", "hi", "ta" etc. Let me write the benchmark script.


Updated task #3 status


Write

#!/usr/bin/env python3
"""
Benchmark Cohere-Transcribe-Indic checkpoints on indic-asr-benchmark-6k dataset.

Usage:
    python benchmark_cohere_transcribe.py \
        --checkpoint /home/ubuntu/training/checkpoints/cohere-transcribe-ckpt-10000 \
        --checkpoint-name ckpt-10000 \
        --batch-size 8

Produces schema-v1 benchmark_outputs (metrics.json, sample_analysis.json, error_analysis.json).
"""

import argparse
import json
import os
import re
import sys
import time
import unicodedata
from collections import Counter, defaultdict
from datetime import datetime, timezone
from pathlib import Path

import numpy as np
import torch
import torchaudio
from datasets import Audio, load_dataset as hf_load_dataset
from jiwer import cer, wer

# Add repo for tokenizer_utils
sys.path.insert(0, "/home/ubuntu/training/cohere-transcribe-indic")
from tokenizer_utils import decode_tokens, load_extended_tokenizer

# ── Constants ────────────────────────────────────────────────────────────

NORMALIZATION_VERSION = "v1"
GPU_NAME = "NVIDIA H200 80GB"
MODEL_ID = "cohere-transcribe"
MODEL_TYPE = "Cohere-Transcribe-Indic-2B"
DATASET = "BayAreaBoys/indic-asr-benchmark-6k"

LANG_CODE_TO_NAME = {
    "as": "Assamese", "bn": "Bengali", "en": "English", "gu": "Gujarati",
    "hi": "Hindi", "kn": "Kannada", "ml": "Malayalam", "mr": "Marathi",
    "or": "Odia", "pa": "Punjabi", "ta": "Tamil", "te": "Telugu",
}

# Map language names to model's language codes
LANG_NAME_TO_CODE = {v.lower(): k for k, v in LANG_CODE_TO_NAME.items()}

LATIN_SCRIPT_LANGS = {"english"}

# Zero-width characters to strip
ZW_CHARS = re.compile(r"[\u200b-\u200f\u202a-\u202e\ufeff\u00ad]")

# Language-aware punctuation sets
BASE_PUNCT = set("!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~")
INDIC_PUNCT = {"\u0964", "\u0965", "\u0970"}
INDIC_REMOVE_PUNCT = {"\u0965", "\u0970"}

QUOTE_MAP = {
    "\u2018": "'", "\u2019": "'", "\u201c": '"', "\u201d": '"',
    "\u2013": "-", "\u2014": "-", "\u2015": "-", "\u2026": "...",
}

EN_NUMBER_WORDS = {
    "zero": "0", "one": "1", "two": "2", "three": "3", "four": "4",
    "five": "5", "six": "6", "seven": "7", "eight": "8", "nine": "9",
    "ten": "10", "eleven": "11", "twelve": "12", "thirteen": "13",
    "fourteen": "14", "fifteen": "15", "sixteen": "16", "seventeen": "17",
    "eighteen": "18", "nineteen": "19", "twenty": "20", "thirty": "30",
    "forty": "40", "fifty": "50", "sixty": "60", "seventy": "70",
    "eighty": "80", "ninety": "90", "hundred": "100", "thousand": "1000",
    "lakh": "100000", "crore": "10000000",
}

HI_NUMBER_WORDS = {
    "शून्य": "0", "एक": "1", "दो": "2", "तीन": "3", "चार": "4",
    "पांच": "5", "छह": "6", "सात": "7", "आठ": "8", "नौ": "9",
    "दस": "10", "ग्यारह": "11", "बारह": "12", "तेरह": "13",
    "चौदह": "14", "पंद्रह": "15", "सोलह": "16", "सत्रह": "17",
    "अठारह": "18", "उन्नीस": "19", "बीस": "20",
    "सौ": "100", "हज़ार": "1000", "हजार": "1000",
    "लाख": "100000", "करोड़": "10000000",
}

# ── Normalization Pipeline ───────────────────────────────────────────────

def norm_raw(text: str) -> str:
    text = unicodedata.normalize("NFC", text)
    return text.strip()

def norm_standard(text: str, language: str = "") -> str:
    text = unicodedata.normalize("NFKC", text)
    text = ZW_CHARS.sub("", text)
    text = re.sub(r"\s+", " ", text).strip()
    for old, new in QUOTE_MAP.items():
        text = text.replace(old, new)
    text = text.replace("\u0965", "\u0964")
    text = "".join(c for c in text if c not in BASE_PUNCT)
    text = "".join(c for c in text if c not in INDIC_PUNCT and c not in INDIC_REMOVE_PUNCT)
    text = text.lower()
    text = re.sub(r"\s+", " ", text).strip()
    return text

def norm_numcanon(text: str, language: str = "") -> str:
    text = norm_standard(text, language)
    text = re.sub(r"(\d),(\d)", r"\1\2", text)
    text = re.sub(r"(\d),(\d)", r"\1\2", text)
    if language == "english":
        words = text.split()
        text = " ".join(EN_NUMBER_WORDS.get(w, w) for w in words)
    if language == "hindi":
        words = text.split()
        text = " ".join(HI_NUMBER_WORDS.get(w, w) for w in words)
    return text

def norm_mer(text: str, language: str = "") -> str:
    return norm_standard(text, language).replace(" ", "")

# ── space_norm_wer computation ───────────────────────────────────────

def space_norm_wer_sample(ref_norm: str, hyp_norm: str):
    ref_words = ref_norm.split()
    if not ref_words:
        return (0, 0)
    ref_nospace = ref_norm.replace(" ", "")
    hyp_nospace = hyp_norm.replace(" ", "")
    total_words = len(ref_words)
    if ref_nospace == hyp_nospace:
        return (0, total_words)

    char_to_word = []
    for word_idx, word in enumerate(ref_words):
        for _ in word:
            char_to_word.append(word_idx)

    n = len(ref_nospace)
    m = len(hyp_nospace)

    dp = [[0] * (m + 1) for _ in range(n + 1)]
    for i in range(n + 1):
        dp[i][0] = i
    for j in range(m + 1):
        dp[0][j] = j
    for i in range(1, n + 1):
        for j in range(1, m + 1):
            if ref_nospace[i - 1] == hyp_nospace[j - 1]:
                dp[i][j] = dp[i - 1][j - 1]
            else:
                dp[i][j] = 1 + min(dp[i - 1][j - 1], dp[i - 1][j], dp[i][j - 1])

    touched = set()
    i, j = n, m
    while i > 0 or j > 0:
        if i > 0 and j > 0 and ref_nospace[i-1] == hyp_nospace[j-1] and dp[i][j] == dp[i-1][j-1]:
            i -= 1; j -= 1
        elif i > 0 and j > 0 and dp[i][j] == dp[i-1][j-1] + 1:
            touched.add(char_to_word[i-1]); i -= 1; j -= 1
        elif i > 0 and dp[i][j] == dp[i-1][j] + 1:
            touched.add(char_to_word[i-1]); i -= 1
        elif j > 0 and dp[i][j] == dp[i][j-1] + 1:
            if i > 0:
                touched.add(char_to_word[i-1])
            elif i < n:
                touched.add(char_to_word[i])
            j -= 1
        else:
            break

    return (len(touched), total_words)

# ── Per-sample WER ───────────────────────────────────────────────────────

def sample_wer(ref: str, hyp: str) -> float:
    if not ref.strip():
        return 0.0 if not hyp.strip() else 100.0
    if not hyp.strip():
        return 100.0
    return round(wer(ref, hyp) * 100, 2)

# ── Flag detection ───────────────────────────────────────────────────────

def detect_flags(ref: str, hyp: str, ref_n: str, hyp_n: str,
                 wer_norm_val: float, mer_val: float) -> list:
    flags = []
    if ref == hyp:
        flags.append("exact_match")
    if ref_n == hyp_n:
        flags.append("exact_match_norm")
    if not hyp.strip():
        flags.append("empty_hypothesis")
    if ref != hyp and ref_n == hyp_n:
        flags.append("punctuation_only_diff")
    ref_digits = re.findall(r"\d+", ref_n)
    hyp_digits = re.findall(r"\d+", hyp_n)
    if ref_digits != hyp_digits and (ref_digits or hyp_digits):
        flags.append("numeric_mismatch")
    if wer_norm_val > 80:
        flags.append("high_wer")
    if wer_norm_val > mer_val and mer_val < 100:
        flags.append("spacing_error")
    return flags

# ── Error Analysis ───────────────────────────────────────────────────────

def compute_error_analysis(samples_by_lang: dict) -> dict:
    analysis = {}

    for lang, samples in sorted(samples_by_lang.items()):
        subs = Counter()
        insertions = Counter()
        deletions = Counter()
        numeric_mismatch_count = 0
        punct_only_count = 0
        spacing_count = 0
        entity_count = 0
        script_confusion_count = 0
        empty_count = 0

        wer_scores = []

        for s in samples:
            ref_words = s["ref_norm"].split()
            hyp_words = s["hyp_norm"].split()
            flags = s.get("flags", [])

            if "empty_hypothesis" in flags: empty_count += 1
            if "numeric_mismatch" in flags: numeric_mismatch_count += 1
            if "punctuation_only_diff" in flags: punct_only_count += 1
            if "script_mismatch" in flags: script_confusion_count += 1
            if "spacing_error" in flags: spacing_count += 1

            wer_scores.append((s["id"], s.get("wer_norm", 100.0)))

            ref_counter = Counter(ref_words)
            hyp_counter = Counter(hyp_words)
            for word, count in (ref_counter - hyp_counter).items():
                deletions[word] += count
            for word, count in (hyp_counter - ref_counter).items():
                insertions[word] += count
            min_len = min(len(ref_words), len(hyp_words))
            for i in range(min_len):
                if ref_words[i] != hyp_words[i]:
                    subs[(ref_words[i], hyp_words[i])] += 1
            for rw, hw in zip(ref_words[:min_len], hyp_words[:min_len]):
                if rw != hw and (rw[0:1].isupper() or any(c.isupper() for c in rw)):
                    entity_count += 1

        wer_scores.sort(key=lambda x: x[1])
        best = [sid for sid, _ in wer_scores[:3]]
        worst = [sid for sid, _ in wer_scores[-3:]]
        numeric_examples = [s["id"] for s in samples if "numeric_mismatch" in s.get("flags", [])][:3]
        entity_examples = [s["id"] for s in samples if "entity_mismatch" in s.get("flags", [])][:3]

        analysis[lang] = {
            "top_substitutions": [
                {"ref": r, "hyp": h, "count": c} for (r, h), c in subs.most_common(20)
            ],
            "top_insertions": [
                {"word": w, "count": c} for w, c in insertions.most_common(20)
            ],
            "top_deletions": [
                {"word": w, "count": c} for w, c in deletions.most_common(20)
            ],
            "error_buckets": {
                "numeric_mismatch_count": numeric_mismatch_count,
                "punctuation_only_count": punct_only_count,
                "spacing_tokenization_count": spacing_count,
                "entity_mismatch_count": entity_count,
                "script_confusion_count": script_confusion_count,
                "empty_hypothesis_count": empty_count,
            },
            "examples": {
                "worst_samples": worst,
                "best_samples": best,
                "numeric_mismatch_samples": numeric_examples,
                "entity_mismatch_samples": entity_examples,
            }
        }

    analysis["__summary__"] = {
        "model_diagnosis": "recognition-limited",
        "primary_error_source": "recognition",
        "numeric_verbalization_impact": "moderate",
        "formatting_impact": "low",
        "worst_languages": [],
        "best_languages": [],
    }

    return analysis

# ── Inference ────────────────────────────────────────────────────────────

def run_inference(model, processor, tokenizer, dataset, batch_size=1, max_new_tokens=256, device="cuda:0"):
    """Run inference sample-by-sample (batched mel extraction, sequential generation)."""
    results = []
    total = len(dataset)

    for start in range(0, total, batch_size):
        end = min(start + batch_size, total)
        batch = dataset.select(range(start, end))

        for i, sample in enumerate(batch):
            lang_code = sample["lang_code"]
            audio = sample["audio"]
            wav = np.array(audio["array"], dtype=np.float32)
            sr = audio["sampling_rate"]

            # Resample if needed
            if sr != 16000:
                wav_t = torch.from_numpy(wav).unsqueeze(0)
                wav_t = torchaudio.functional.resample(wav_t, sr, 16000)
                wav = wav_t.squeeze(0).numpy()

            # Extract mel features
            features = processor.feature_extractor(
                wav, sampling_rate=16000, return_tensors="pt"
            )
            input_features = features["input_features"].to(device, dtype=torch.bfloat16)
            length = features["length"].to(device)

            # Build decoder prompt
            prompt_str = model.build_prompt(language=lang_code, punctuation=True)
            prompt_ids = tokenizer.encode(prompt_str, add_special_tokens=False)
            decoder_input_ids = torch.tensor([prompt_ids], device=device)

            # Generate
            with torch.no_grad():
                try:
                    outputs = model.generate(
                        input_features=input_features,
                        length=length,
                        decoder_input_ids=decoder_input_ids,
                        max_new_tokens=max_new_tokens,
                    )
                    hyp = decode_tokens(tokenizer, outputs[0]).strip()
                except Exception as e:
                    print(f"  Sample {sample['id']} failed: {e}")
                    hyp = ""

            results.append({
                "id": sample["id"],
                "language": sample["language"],
                "lang_code": lang_code,
                "source": sample["source"],
                "duration": sample["duration"],
                "reference": sample["reference"],
                "hypothesis": hyp,
                "detected_language": lang_code,
            })

        elapsed_pct = end / total * 100
        print(f"  [{end}/{total}] ({elapsed_pct:.0f}%) processed", flush=True)

    return results

# ── Main ─────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="Benchmark Cohere-Transcribe on Indic ASR dataset")
    parser.add_argument("--checkpoint", required=True, help="Path to model checkpoint directory")
    parser.add_argument("--checkpoint-name", default=None)
    parser.add_argument("--hf-dataset", default="BayAreaBoys/indic-asr-benchmark-6k")
    parser.add_argument("--cache-dir", default="/home/ubuntu/training/datasets/indic-asr-benchmark-6k")
    parser.add_argument("--output-dir", default="/home/ubuntu/training/benchmark_outputs/cohere-transcribe")
    parser.add_argument("--batch-size", type=int, default=1)
    parser.add_argument("--max-new-tokens", type=int, default=256)
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument("--languages", nargs="+", default=None)
    args = parser.parse_args()

    ckpt_name = args.checkpoint_name or Path(args.checkpoint).name
    output_dir = Path(args.output_dir) / ckpt_name
    output_dir.mkdir(parents=True, exist_ok=True)

    # Load model
    from transformers import AutoModelForSpeechSeq2Seq, AutoFeatureExtractor

    print(f"Loading checkpoint: {args.checkpoint}")
    t0 = time.time()

    # Load feature extractor from checkpoint (has custom processing code)
    feature_extractor = AutoFeatureExtractor.from_pretrained(
        args.checkpoint, trust_remote_code=True
    )
    # Move filterbank to GPU for faster feature extraction
    feature_extractor._device = args.device

    # Load extended tokenizer
    tokenizer = load_extended_tokenizer(args.checkpoint)

    # Create a simple processor-like namespace
    class SimpleProcessor:
        pass
    processor = SimpleProcessor()
    processor.feature_extractor = feature_extractor
    processor.tokenizer = tokenizer

    # Load model
    model = AutoModelForSpeechSeq2Seq.from_pretrained(
        args.checkpoint,
        trust_remote_code=True,
        torch_dtype=torch.bfloat16,
    ).to(args.device)
    model.eval()

    load_time = time.time() - t0
    print(f"Model loaded in {load_time:.1f}s")

    # Load dataset
    hf_token = os.environ.get("HF_TOKEN")
    print(f"Loading dataset: {args.hf_dataset}")
    dataset = hf_load_dataset(
        args.hf_dataset, token=hf_token, cache_dir=args.cache_dir, split="train"
    )
    dataset = dataset.cast_column("audio", Audio(sampling_rate=16000, decode=True))

    if args.languages:
        dataset = dataset.filter(lambda x: x["lang_code"] in args.languages)
        print(f"Filtered to {len(dataset)} samples for languages: {args.languages}")

    print(f"Dataset: {len(dataset)} samples")

    # Run inference
    print(f"\nRunning inference (batch_size={args.batch_size})...")
    t0 = time.time()
    results = run_inference(
        model, processor, tokenizer, dataset,
        batch_size=args.batch_size,
        max_new_tokens=args.max_new_tokens,
        device=args.device,
    )
    inference_time = time.time() - t0
    total_audio_sec = sum(r["duration"] for r in results)
    rtf = inference_time / total_audio_sec if total_audio_sec > 0 else 0
    print(f"Inference done in {inference_time:.1f}s for {total_audio_sec:.0f}s audio (RTF={rtf:.4f})")

    # Save raw predictions
    preds_file = output_dir / "predictions.json"
    with open(preds_file, "w", encoding="utf-8") as f:
        json.dump(results, f, ensure_ascii=False, indent=2)
    print(f"Predictions saved: {preds_file}")

    # ── Compute all metric tiers ──
    print("\nComputing metrics...")
    enriched = []
    for p in results:
        lang = p["language"]
        ref = p["reference"]
        hyp = p.get("hypothesis", "")

        ref_raw = norm_raw(ref)
        hyp_raw = norm_raw(hyp)
        ref_n = norm_standard(ref, lang)
        hyp_n = norm_standard(hyp, lang)
        ref_nc = norm_numcanon(ref, lang)
        hyp_nc = norm_numcanon(hyp, lang)
        ref_m = norm_mer(ref, lang)
        hyp_m = norm_mer(hyp, lang)

        wer_raw_val = sample_wer(ref_raw, hyp_raw)
        wer_norm_val = sample_wer(ref_n, hyp_n)

        if not ref_m:
            mer_val = 0.0 if not hyp_m else 100.0
        elif not hyp_m:
            mer_val = 100.0
        else:
            mer_val = round(cer(ref_m, hyp_m) * 100, 2)

        snw_err, snw_total = space_norm_wer_sample(ref_n, hyp_n)
        snw_val = round(snw_err / snw_total * 100, 2) if snw_total > 0 else 0.0

        flags = detect_flags(ref_raw, hyp_raw, ref_n, hyp_n, wer_norm_val, mer_val)

        enriched.append({
            "id": p["id"],
            "language": lang,
            "reference": ref,
            "hypothesis": hyp,
            "ref_raw": ref_raw, "hyp_raw": hyp_raw,
            "ref_norm": ref_n, "hyp_norm": hyp_n,
            "ref_numcanon": ref_nc, "hyp_numcanon": hyp_nc,
            "ref_mer": ref_m, "hyp_mer": hyp_m,
            "detected_language": p.get("detected_language", ""),
            "wer_raw": wer_raw_val, "wer_norm": wer_norm_val,
            "space_norm_wer": snw_val, "snw_err": snw_err, "snw_total": snw_total,
            "mer": mer_val, "flags": flags,
        })

    # ── Per-language metrics ──
    by_lang = defaultdict(list)
    for s in enriched:
        by_lang[s["language"]].append(s)

    metrics = {}
    all_ref_raw, all_hyp_raw = [], []
    all_ref_n, all_hyp_n = [], []
    all_ref_nc, all_hyp_nc = [], []
    all_ref_m, all_hyp_m = [], []
    all_snw_err, all_snw_total = 0, 0

    for lang in sorted(by_lang.keys()):
        samples = by_lang[lang]
        n = len(samples)

        refs_raw = [s["ref_raw"] for s in samples if s["ref_norm"]]
        hyps_raw = [s["hyp_raw"] for s in samples if s["ref_norm"]]
        refs_n = [s["ref_norm"] for s in samples if s["ref_norm"]]
        hyps_n = [s["hyp_norm"] for s in samples if s["ref_norm"]]
        refs_nc = [s["ref_numcanon"] for s in samples if s["ref_norm"]]
        hyps_nc = [s["hyp_numcanon"] for s in samples if s["ref_norm"]]
        refs_m = [s["ref_mer"] for s in samples if s["ref_norm"]]
        hyps_m = [s["hyp_mer"] for s in samples if s["ref_norm"]]

        w_raw = round(wer(refs_raw, hyps_raw) * 100, 2) if refs_raw else 100.0
        w_norm = round(wer(refs_n, hyps_n) * 100, 2) if refs_n else 100.0
        w_nc = round(wer(refs_nc, hyps_nc) * 100, 2) if refs_nc else 100.0
        w_mer = round(cer(refs_m, hyps_m) * 100, 2) if refs_m else 100.0
        c_norm = round(cer(refs_n, hyps_n) * 100, 2) if refs_n else 100.0
        empty = sum(1 for s in samples if "empty_hypothesis" in s["flags"])

        lang_snw_err = sum(s["snw_err"] for s in samples if s["ref_norm"])
        lang_snw_total = sum(s["snw_total"] for s in samples if s["ref_norm"])
        w_snw = round(lang_snw_err / lang_snw_total * 100, 2) if lang_snw_total > 0 else 100.0

        metrics[lang] = {
            "n_samples": n,
            "wer_raw": w_raw, "wer_norm": w_norm, "wer_numcanon": w_nc,
            "space_norm_wer": w_snw, "mer": w_mer, "cer_norm": c_norm,
            "empty_hypotheses": empty,
            "normalization_delta": {
                "raw_to_norm": round(w_raw - w_norm, 2),
                "norm_to_numcanon": round(w_norm - w_nc, 2),
                "norm_to_space_norm": round(w_norm - w_snw, 2),
                "norm_to_mer": round(w_norm - w_mer, 2),
            }
        }

        all_ref_raw.extend(refs_raw); all_hyp_raw.extend(hyps_raw)
        all_ref_n.extend(refs_n); all_hyp_n.extend(hyps_n)
        all_ref_nc.extend(refs_nc); all_hyp_nc.extend(hyps_nc)
        all_ref_m.extend(refs_m); all_hyp_m.extend(hyps_m)
        all_snw_err += lang_snw_err; all_snw_total += lang_snw_total

    # Aggregates
    metrics["__overall__"] = {
        "n_samples": len(all_ref_raw),
        "wer_raw": round(wer(all_ref_raw, all_hyp_raw) * 100, 2),
        "wer_norm": round(wer(all_ref_n, all_hyp_n) * 100, 2),
        "wer_numcanon": round(wer(all_ref_nc, all_hyp_nc) * 100, 2),
        "space_norm_wer": round(all_snw_err / all_snw_total * 100, 2) if all_snw_total > 0 else 100.0,
        "mer": round(cer(all_ref_m, all_hyp_m) * 100, 2),
        "cer_norm": round(cer(all_ref_n, all_hyp_n) * 100, 2),
    }

    lang_keys = sorted(k for k in metrics if not k.startswith("__"))
    metrics["__macro_avg__"] = {
        "n_languages": len(lang_keys),
        "wer_raw": round(np.mean([metrics[k]["wer_raw"] for k in lang_keys]), 2),
        "wer_norm": round(np.mean([metrics[k]["wer_norm"] for k in lang_keys]), 2),
        "wer_numcanon": round(np.mean([metrics[k]["wer_numcanon"] for k in lang_keys]), 2),
        "space_norm_wer": round(np.mean([metrics[k]["space_norm_wer"] for k in lang_keys]), 2),
        "mer": round(np.mean([metrics[k]["mer"] for k in lang_keys]), 2),
        "cer_norm": round(np.mean([metrics[k]["cer_norm"] for k in lang_keys]), 2),
    }

    import importlib.metadata
    jiwer_ver = importlib.metadata.version("jiwer")
    metrics["__meta__"] = {
        "checkpoint": args.checkpoint,
        "checkpoint_name": ckpt_name,
        "model_id": MODEL_ID,
        "model_type": MODEL_TYPE,
        "dataset": DATASET,
        "batch_size": args.batch_size,
        "inference_time_sec": round(inference_time, 2),
        "total_audio_sec": round(total_audio_sec, 2),
        "rtf": round(rtf, 4),
        "timestamp": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "gpu": GPU_NAME,
        "framework": "transformers",
        "normalization_version": NORMALIZATION_VERSION,
        "jiwer_version": jiwer_ver,
    }

    # ── sample_analysis.json ──
    sample_analysis = []
    for s in enriched:
        sample_analysis.append({
            "id": s["id"], "language": s["language"],
            "reference": s["reference"], "hypothesis": s["hypothesis"],
            "ref_norm": s["ref_norm"], "hyp_norm": s["hyp_norm"],
            "ref_numcanon": s["ref_numcanon"], "hyp_numcanon": s["hyp_numcanon"],
            "ref_mer": s["ref_mer"], "hyp_mer": s["hyp_mer"],
            "detected_language": s["detected_language"],
            "wer_raw": s["wer_raw"], "wer_norm": s["wer_norm"],
            "space_norm_wer": s["space_norm_wer"], "mer": s["mer"],
            "flags": s["flags"],
        })

    # ── error_analysis.json ──
    samples_for_errors = defaultdict(list)
    for s in sample_analysis:
        samples_for_errors[s["language"]].append(s)

    error_analysis = compute_error_analysis(samples_for_errors)

    sorted_langs = sorted(lang_keys, key=lambda k: metrics[k]["wer_norm"])
    error_analysis["__summary__"]["best_languages"] = sorted_langs[:3]
    error_analysis["__summary__"]["worst_languages"] = sorted_langs[-3:]

    avg_raw_to_norm = np.mean([metrics[k]["normalization_delta"]["raw_to_norm"] for k in lang_keys])
    avg_norm_to_nc = np.mean([metrics[k]["normalization_delta"]["norm_to_numcanon"] for k in lang_keys])

    if avg_raw_to_norm > 10:
        error_analysis["__summary__"]["model_diagnosis"] = "formatting-limited"
        error_analysis["__summary__"]["formatting_impact"] = "high"
    elif avg_norm_to_nc > 5:
        error_analysis["__summary__"]["model_diagnosis"] = "numeric-limited"
        error_analysis["__summary__"]["numeric_verbalization_impact"] = "high"
    else:
        error_analysis["__summary__"]["model_diagnosis"] = "recognition-limited"
        error_analysis["__summary__"]["formatting_impact"] = "low" if avg_raw_to_norm < 5 else "moderate"
        error_analysis["__summary__"]["numeric_verbalization_impact"] = "low" if avg_norm_to_nc < 1 else "moderate"

    # ── Write output files ──
    with open(output_dir / "metrics.json", "w", encoding="utf-8") as f:
        json.dump(metrics, f, indent=2, ensure_ascii=False)

    with open(output_dir / "sample_analysis.json", "w", encoding="utf-8") as f:
        json.dump(sample_analysis, f, indent=2, ensure_ascii=False)

    with open(output_dir / "error_analysis.json", "w", encoding="utf-8") as f:
        json.dump(error_analysis, f, indent=2, ensure_ascii=False)

    # Print summary
    ov = metrics["__overall__"]
    print(f"\n{'='*85}")
    print(f"  BENCHMARK RESULTS: {ckpt_name} ({MODEL_ID})")
    print(f"{'='*85}")
    print(f"{'Language':<14} {'N':>5} {'WER-raw':>8} {'WER-n':>8} {'WER-nc':>8} {'SNW':>8} {'MER':>8} {'CER-n':>8}")
    print(f"{'-'*14} {'-'*5} {'-'*8} {'-'*8} {'-'*8} {'-'*8} {'-'*8} {'-'*8}")

    for lang in sorted(k for k in metrics if not k.startswith("__")):
        m = metrics[lang]
        print(f"{lang:<14} {m['n_samples']:>5} {m['wer_raw']:>8.2f} {m['wer_norm']:>8.2f} "
              f"{m['wer_numcanon']:>8.2f} {m['space_norm_wer']:>8.2f} {m['mer']:>8.2f} {m['cer_norm']:>8.2f}")

    print(f"{'-'*14} {'-'*5} {'-'*8} {'-'*8} {'-'*8} {'-'*8} {'-'*8} {'-'*8}")
    print(f"{'OVERALL':<14} {ov['n_samples']:>5} {ov['wer_raw']:>8.2f} {ov['wer_norm']:>8.2f} "
          f"{ov['wer_numcanon']:>8.2f} {ov['space_norm_wer']:>8.2f} {ov['mer']:>8.2f} {ov['cer_norm']:>8.2f}")
    ma = metrics["__macro_avg__"]
    print(f"{'MACRO-AVG':<14} {ma['n_languages']:>5} {ma['wer_raw']:>8.2f} {ma['wer_norm']:>8.2f} "
          f"{ma['wer_numcanon']:>8.2f} {ma['space_norm_wer']:>8.2f} {ma['mer']:>8.2f} {ma['cer_norm']:>8.2f}")
    print(f"{'='*85}")
    print(f"\nOutputs saved to: {output_dir}")
    print(f"  metrics.json, sample_analysis.json, error_analysis.json")

if __name__ == "__main__":
    main()