Here's the clean apples-to-apples picture for **ckpt-70000**:

### Clean Metrics (excluding language mismatches + numerics)

| Language | N | WER_raw | WER_norm | CER_norm | Delta (raw-norm) |
|----------|---|---------|----------|----------|-----------------|
| Hindi | 496 | 11.85 | **11.85** | 3.76 | 0.0 |
| Assamese | 499 | 14.26 | **14.07** | 6.05 | +0.19 |
| Gujarati | 495 | 16.75 | **16.75** | 5.00 | 0.0 |
| Bengali | 495 | 18.22 | **17.69** | 3.93 | +0.53 |
| Punjabi | 473 | 20.55 | **20.05** | 9.18 | +0.50 |
| Marathi | 499 | 22.30 | **22.30** | 6.24 | 0.0 |
| Telugu | 474 | 25.89 | **25.87** | 5.97 | +0.02 |
| Odia | 456 | 26.53 | **26.30** | 8.09 | +0.23 |
| Tamil | 478 | 29.24 | **29.22** | 6.86 | +0.02 |
| Kannada | 455 | 32.91 | **32.91** | 13.09 | 0.0 |
| English | 338 | 19.36 | **8.48** | 4.18 | +10.88 |
| Malayalam | 498 | 38.92 | **38.92** | 7.49 | 0.0 |
| **OVERALL** | 5656 | 22.68 | **21.94** | 6.84 | +0.74 |
| **Indic only** | 5318 | 22.88 | **22.75** | 6.97 | +0.13 |

### Key findings:

**WER_raw ≈ WER_norm for all Indic languages** — delta is 0.0-0.5 points. The model outputs virtually no punctuation or casing for Indic scripts. The only language with a formatting gap is English (+10.88) due to casing.

**Biggest improvements from filtering:**
- **Kannada**: 38.84 → 32.91 (-5.93) — 44 language mismatches were script-confused (outputting Tamil/Gujarati instead of Kannada)
- **Odia**: 33.35 → 26.30 (-7.05) — 42 mismatches (outputting Devanagari/Bengali instead of Odia)
- **English**: 23.92 → 8.48 (-15.44) — 80 samples where model output Indic script for English audio + 97 numeric samples
- **Punjabi**: 23.90 → 20.05 (-3.85) — 27 samples output in Devanagari instead of Gurmukhi

**CER_norm dropped from 10.32 → 6.84** — nearly half the character errors were from script confusion or numeric handling.

**Bottom line:** The model's *pure recognition* ability on clean Indic text is **22.75% WER_norm** — much better than the headline 25.03%. The main issues are script confusion (Kannada/Odia/Punjabi/Telugu) and numeric handling, not transcription quality.


i noticed a lot of language mismatches too. i think, can we just remove the samples where the originasl and generated transcriptions's language is mismatched. just wanna see straight up apples-apples…


can you fix the --backend vllm route. i can't keep waiting so long for using teh transformers implementation and that is very slow. check what the issue is and fix it.


1	#!/usr/bin/env python3
2	"""
3	Benchmark Gemma3n-E2B-ASR checkpoints on indic-asr-benchmark-6k dataset.
4
5	Usage:
6	    python benchmark_gemma3n_asr.py \
7	        --checkpoint /home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-10000 \
8	        --checkpoint-name ckpt-10000 \
9	        --output-dir /home/ubuntu/training/benchmark_results_gemma3n \
10	        --batch-size 1
11
12	Metrics computed per language and overall:
13	    - WER  (Word Error Rate)
14	    - CER  (Character Error Rate)
15	    - WER-norm (WER after unicode/punctuation/space normalization)
16	    - CER-norm (CER after normalization)
17	"""
18
19	import argparse
20	import json
21	import os
22	import re
23	import sys
24	import time
25	import unicodedata
26	from collections import defaultdict
27	from pathlib import Path
28
29	import numpy as np
30	import torch
31	import torchaudio
32	from datasets import load_dataset as hf_load_dataset
33	from jiwer import cer, wer
34	from transformers import AutoModelForImageTextToText, AutoProcessor
35
36	# Language code → display name mapping
37	LANG_CODE_TO_NAME = {
38	    "as": "Assamese", "bn": "Bengali", "en": "English", "gu": "Gujarati",
39	    "hi": "Hindi", "kn": "Kannada", "ml": "Malayalam", "mr": "Marathi",
40	    "or": "Odia", "pa": "Punjabi", "ta": "Tamil", "te": "Telugu",
41	}
42
43
44	# ── Text normalization (standard ASR eval practice) ──────────────────────
45
46	def normalize_text(text: str) -> str:
47	    """Normalize text for fair WER/CER comparison.
48	    Unicode NFC → lowercase → remove punctuation → collapse whitespace → strip.
49	    """
50	    text = unicodedata.normalize("NFC", text)
51	    text = text.lower()
52	    text = re.sub(r"[^\w\s]", "", text, flags=re.UNICODE)
53	    text = re.sub(r"\s+", " ", text).strip()
54	    return text
55
56
57	def load_audio_array(sample, target_sr=16000):
58	    """Extract audio as numpy float32 array at target sample rate."""
59	    audio = sample["audio"]
60	    wav = np.array(audio["array"], dtype=np.float32)
61	    sr = audio["sampling_rate"]
62
63	    if sr != target_sr:
64	        wav_tensor = torch.from_numpy(wav).unsqueeze(0)
65	        wav_tensor = torchaudio.functional.resample(wav_tensor, sr, target_sr)
66	        wav = wav_tensor.squeeze(0).numpy()
67
68	    return wav
69
70
71	def run_inference(model, processor, dataset, batch_size: int = 16, max_new_tokens: int = 256):
72	    """Run batched inference on HF dataset and return list of result dicts."""
73	    results = []
74	    total = len(dataset)
75
76	    for start in range(0, total, batch_size):
77	        end = min(start + batch_size, total)
78	        batch = dataset.select(range(start, end))
79
80	        # Prepare batch: audio arrays + prompts
81	        wavs = []
82	        prompts = []
83	        batch_meta = []
84	        for sample in batch:
85	            lang_name = LANG_CODE_TO_NAME.get(sample["lang_code"], sample["language"].capitalize())
86	            user_content = f"<audio_soft_token>\nTranscribe the audio.\nLanguage: {lang_name}"
87	            messages = [{"role": "user", "content": user_content}]
88
89	            wavs.append(load_audio_array(sample))
90	            prompts.append(processor.apply_chat_template(
91	                messages, add_generation_prompt=True, tokenize=False
92	            ))
93	            batch_meta.append({
94	                "id": sample["id"],
95	                "language": sample["language"],
96	                "lang_code": sample["lang_code"],
97	                "source": sample["source"],
98	                "duration": sample["duration"],
99	                "reference": sample["reference"],
100	            })
101
102	        try:
103	            # Process batch
104	            inputs = processor(
105	                text=prompts, audio=wavs, sampling_rate=16000,
106	                return_tensors="pt", padding=True,
107	            )
108	            inputs = {k: v.to(model.device) if isinstance(v, torch.Tensor) else v
109	                      for k, v in inputs.items()}
110	            input_len = inputs["input_ids"].shape[1]
111
112	            # Generate
113	            with torch.no_grad():
114	                output_ids = model.generate(
115	                    **inputs, max_new_tokens=max_new_tokens, do_sample=False,
116	                )
117
118	            # Decode each sample in the batch
119	            for i in range(len(batch_meta)):
120	                gen_ids = output_ids[i][input_len:]
121	                hypothesis = processor.tokenizer.decode(gen_ids, skip_special_tokens=True).strip()
122	                results.append({**batch_meta[i], "hypothesis": hypothesis})
123
124	        except Exception as e:
125	            print(f"  Batch {start}-{end} FAILED ({e}), falling back to sequential...", flush=True)
126	            for i, sample in enumerate(batch):
127	                try:
128	                    inp = processor(
129	                        text=[prompts[i]], audio=[wavs[i]], sampling_rate=16000,
130	                        return_tensors="pt",
131	                    )
132	                    inp = {k: v.to(model.device) if isinstance(v, torch.Tensor) else v
133	                           for k, v in inp.items()}
134	                    with torch.no_grad():
135	                        out = model.generate(**inp, max_new_tokens=max_new_tokens, do_sample=False)
136	                    gen_ids = out[0][inp["input_ids"].shape[1]:]
137	                    hypothesis = processor.tokenizer.decode(gen_ids, skip_special_tokens=True).strip()
138	                except Exception as e2:
139	                    print(f"  Sample {start+i} ({batch_meta[i]['id']}) FAILED: {e2}")
140	                    hypothesis = ""
141	                results.append({**batch_meta[i], "hypothesis": hypothesis})
142
143	        elapsed_pct = end / total * 100
144	        print(f"  [{end}/{total}] ({elapsed_pct:.0f}%) processed", flush=True)
145
146	    return results
147
148
149	def compute_metrics(results: list) -> dict:
150	    """Compute WER, CER, normalized-WER/CER per language and overall."""
151	    by_lang = defaultdict(lambda: {"refs": [], "hyps": [], "refs_norm": [], "hyps_norm": []})
152
153	    for r in results:
154	        lang = r["language"]
155	        ref = r["reference"]
156	        hyp = r.get("hypothesis", "")
157
158	        ref_norm = normalize_text(ref)
159	        hyp_norm = normalize_text(hyp)
160
161	        if not ref_norm:
162	            continue
163
164	        by_lang[lang]["refs"].append(ref)
165	        by_lang[lang]["hyps"].append(hyp)
166	        by_lang[lang]["refs_norm"].append(ref_norm)
167	        by_lang[lang]["hyps_norm"].append(hyp_norm)
168
169	    metrics = {}
170	    all_refs, all_hyps = [], []
171	    all_refs_norm, all_hyps_norm = [], []
172
173	    for lang in sorted(by_lang.keys()):
174	        d = by_lang[lang]
175	        n = len(d["refs"])
176	        lang_wer = wer(d["refs"], d["hyps"])
177	        lang_cer = cer(d["refs"], d["hyps"])
178	        lang_wer_norm = wer(d["refs_norm"], d["hyps_norm"])
179	        lang_cer_norm = cer(d["refs_norm"], d["hyps_norm"])
180
181	        empty_count = sum(1 for h in d["hyps"] if not h.strip())
182
183	        metrics[lang] = {
184	            "n_samples": n,
185	            "wer": round(lang_wer * 100, 2),
186	            "cer": round(lang_cer * 100, 2),
187	            "wer_normalized": round(lang_wer_norm * 100, 2),
188	            "cer_normalized": round(lang_cer_norm * 100, 2),
189	            "empty_hypotheses": empty_count,
190	        }
191
192	        all_refs.extend(d["refs"])
193	        all_hyps.extend(d["hyps"])
194	        all_refs_norm.extend(d["refs_norm"])
195	        all_hyps_norm.extend(d["hyps_norm"])
196
197	    # Overall (corpus-level)
198	    metrics["__overall__"] = {
199	        "n_samples": len(all_refs),
200	        "wer": round(wer(all_refs, all_hyps) * 100, 2),
201	        "cer": round(cer(all_refs, all_hyps) * 100, 2),
202	        "wer_normalized": round(wer(all_refs_norm, all_hyps_norm) * 100, 2),
203	        "cer_normalized": round(cer(all_refs_norm, all_hyps_norm) * 100, 2),
204	    }
205
206	    # Macro average across languages
207	    lang_keys = [k for k in metrics if not k.startswith("__")]
208	    metrics["__macro_avg__"] = {
209	        "n_languages": len(lang_keys),
210	        "wer": round(np.mean([metrics[k]["wer"] for k in lang_keys]), 2),
211	        "cer": round(np.mean([metrics[k]["cer"] for k in lang_keys]), 2),
212	        "wer_normalized": round(np.mean([metrics[k]["wer_normalized"] for k in lang_keys]), 2),
213	        "cer_normalized": round(np.mean([metrics[k]["cer_normalized"] for k in lang_keys]), 2),
214	    }
215
216	    return metrics
217
218
219	def print_metrics_table(metrics: dict, checkpoint_name: str):
220	    """Pretty print metrics as a table."""
221	    print(f"\n{'='*85}")
222	    print(f"  BENCHMARK RESULTS: {checkpoint_name}")
223	    print(f"{'='*85}")
224	    print(f"{'Language':<14} {'N':>5} {'WER%':>8} {'CER%':>8} {'WER-n%':>8} {'CER-n%':>8}")
225	    print(f"{'-'*14} {'-'*5} {'-'*8} {'-'*8} {'-'*8} {'-'*8}")
226
227	    for lang in sorted(k for k in metrics if not k.startswith("__")):
228	        m = metrics[lang]
229	        print(f"{lang:<14} {m['n_samples']:>5} {m['wer']:>8.2f} {m['cer']:>8.2f} "
230	              f"{m['wer_normalized']:>8.2f} {m['cer_normalized']:>8.2f}")
231
232	    print(f"{'-'*14} {'-'*5} {'-'*8} {'-'*8} {'-'*8} {'-'*8}")
233
234	    ov = metrics["__overall__"]
235	    print(f"{'OVERALL':<14} {ov['n_samples']:>5} {ov['wer']:>8.2f} {ov['cer']:>8.2f} "
236	          f"{ov['wer_normalized']:>8.2f} {ov['cer_normalized']:>8.2f}")
237
238	    ma = metrics["__macro_avg__"]
239	    print(f"{'MACRO-AVG':<14} {ma['n_languages']:>5} {ma['wer']:>8.2f} {ma['cer']:>8.2f} "
240	          f"{ma['wer_normalized']:>8.2f} {ma['cer_normalized']:>8.2f}")
241	    print(f"{'='*85}\n")
242
243
244	def save_sample_analysis(results: list, output_dir: Path, n_per_lang: int = 5):
245	    """Save a few sample predictions per language for manual inspection."""
246	    by_lang = defaultdict(list)
247	    for r in results:
248	        by_lang[r["language"]].append(r)
249
250	    samples = []
251	    for lang in sorted(by_lang.keys()):
252	        lang_results = by_lang[lang]
253	        for r in lang_results[:n_per_lang]:
254	            samples.append({
255	                "id": r["id"],
256	                "language": r["language"],
257	                "reference": r["reference"],
258	                "hypothesis": r["hypothesis"],
259	                "ref_normalized": normalize_text(r["reference"]),
260	                "hyp_normalized": normalize_text(r["hypothesis"]),
261	            })
262
263	    path = output_dir / "sample_analysis.json"
264	    with open(path, "w", encoding="utf-8") as f:
265	        json.dump(samples, f, ensure_ascii=False, indent=2)
266	    print(f"Sample analysis saved: {path}")
267
268
269	def main():
270	    parser = argparse.ArgumentParser(description="Benchmark Gemma3n-E2B-ASR on Indic ASR dataset")
271	    parser.add_argument("--checkpoint", required=True, help="Path to model checkpoint directory")
272	    parser.add_argument("--checkpoint-name", default=None,
273	                        help="Name for this checkpoint in output files (default: dirname)")
274	    parser.add_argument("--hf-dataset", default="BayAreaBoys/indic-asr-benchmark-6k",
275	                        help="HuggingFace dataset name")
276	    parser.add_argument("--hf-token", default=None, help="HuggingFace token (or set HF_TOKEN env)")
277	    parser.add_argument("--cache-dir", default="/home/ubuntu/training/datasets/indic-asr-benchmark-6k")
278	    parser.add_argument("--output-dir", default="/home/ubuntu/training/benchmark_results_gemma3n")
279	    parser.add_argument("--batch-size", type=int, default=16,
280	                        help="Samples per batch for offline batched inference")
281	    parser.add_argument("--max-new-tokens", type=int, default=256)
282	    parser.add_argument("--device", default="cuda:0")
283	    parser.add_argument("--languages", nargs="+", default=None,
284	                        help="Filter to specific language codes (e.g. hi en ta)")
285	    args = parser.parse_args()
286
287	    ckpt_name = args.checkpoint_name or Path(args.checkpoint).name
288	    output_dir = Path(args.output_dir) / ckpt_name
289	    output_dir.mkdir(parents=True, exist_ok=True)
290
291	    # Load model and processor
292	    print(f"Loading checkpoint: {args.checkpoint}")
293	    t0 = time.time()
294
295	    # Load tokenizer first, patch missing audio_token_id (version mismatch workaround)
296	    from transformers import AutoTokenizer, AutoConfig
297	    config = AutoConfig.from_pretrained(args.checkpoint, trust_remote_code=True)
298	    tokenizer = AutoTokenizer.from_pretrained(args.checkpoint, trust_remote_code=True, padding_side="left")
299
300	    # Gemma3nProcessor expects special token attrs on the tokenizer (added in transformers 5.x)
301	    # Patch them from config for compatibility with transformers 4.x
302	    _token_patches = {
303	        "audio_token_id": getattr(config, "audio_token_id", 262273),
304	        "audio_token": "<audio_soft_token>",
305	        "boa_token_id": getattr(config, "boa_token_id", 256000),
306	        "boa_token": "<start_of_audio>",
307	        "eoa_token_id": getattr(config, "eoa_token_id", 262272),
308	        "eoa_token": "<end_of_audio>",
309	        "boi_token_id": getattr(config, "boi_token_id", 255999),
310	        "boi_token": "<start_of_image>",
311	        "eoi_token_id": getattr(config, "eoi_token_id", 262144),
312	        "eoi_token": "<end_of_image>",
313	        "image_token_id": getattr(config, "image_token_id", 262145),
314	        "image_token": "<image_soft_token>",
315	    }
316	    for attr, val in _token_patches.items():
317	        if not hasattr(tokenizer, attr):
318	            setattr(tokenizer, attr, val)
319
320	    processor = AutoProcessor.from_pretrained(
321	        args.checkpoint, trust_remote_code=True, tokenizer=tokenizer
322	    )
323	    model = AutoModelForImageTextToText.from_pretrained(
324	        args.checkpoint,
325	        torch_dtype=torch.bfloat16,
326	        device_map=args.device,
327	        trust_remote_code=True,
328	        attn_implementation="sdpa",
329	    )
330	    model.eval()
331	    # Enable KV cache for faster autoregressive generation (training disables it)
332	    model.config.use_cache = True
333	    if hasattr(model.config, "text_config"):
334	        model.config.text_config.use_cache = True
335
336	    load_time = time.time() - t0
337	    print(f"Model loaded in {load_time:.1f}s")
338
339	    # Load dataset
340	    hf_token = args.hf_token or os.environ.get("HF_TOKEN")
341	    print(f"Loading dataset: {args.hf_dataset}")
342	    dataset = hf_load_dataset(
343	        args.hf_dataset, token=hf_token, cache_dir=args.cache_dir, split="train"
344	    )
345
346	    # Filter languages if specified
347	    if args.languages:
348	        dataset = dataset.filter(lambda x: x["lang_code"] in args.languages)
349	        print(f"Filtered to {len(dataset)} samples for languages: {args.languages}")
350
351	    print(f"Dataset: {len(dataset)} samples")
352
353	    # Run inference
354	    print(f"\nRunning inference (batch_size={args.batch_size})...")
355	    t0 = time.time()
356	    results = run_inference(model, processor, dataset, args.batch_size, args.max_new_tokens)
357	    inference_time = time.time() - t0
358
359	    total_audio_sec = sum(r["duration"] for r in results)
360	    rtf = inference_time / total_audio_sec if total_audio_sec > 0 else 0
361	    print(f"Inference done in {inference_time:.1f}s for {total_audio_sec:.0f}s audio (RTF={rtf:.4f})")
362
363	    # Save raw predictions
364	    preds_file = output_dir / "predictions.json"
365	    with open(preds_file, "w", encoding="utf-8") as f:
366	        json.dump(results, f, ensure_ascii=False, indent=2)
367	    print(f"Predictions saved: {preds_file}")
368
369	    # Compute metrics
370	    metrics = compute_metrics(results)
371	    metrics["__meta__"] = {
372	        "checkpoint": args.checkpoint,
373	        "checkpoint_name": ckpt_name,
374	        "model_type": "gemma3n-E2B-asr",
375	        "dataset": args.hf_dataset,
376	        "batch_size": args.batch_size,
377	        "inference_time_sec": round(inference_time, 2),
378	        "total_audio_sec": round(total_audio_sec, 2),
379	        "rtf": round(rtf, 4),
380	    }
381
382	    metrics_file = output_dir / "metrics.json"
383	    with open(metrics_file, "w") as f:
384	        json.dump(metrics, f, indent=2, ensure_ascii=False)
385	    print(f"Metrics saved: {metrics_file}")
386
387	    # Print table
388	    print_metrics_table(metrics, ckpt_name)
389
390	    # Save sample analysis
391	    save_sample_analysis(results, output_dir)
392
393
394	if __name__ == "__main__":
395	    main()
396

#!/usr/bin/env python3
"""
Benchmark Gemma3n-E2B-ASR checkpoints on indic-asr-benchmark-6k dataset.

Usage:
    python benchmark_gemma3n_asr.py \
        --checkpoint /home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-10000 \
        --checkpoint-name ckpt-10000 \
        --output-dir /home/ubuntu/training/benchmark_results_gemma3n \
        --batch-size 1

Metrics computed per language and overall:
    - WER  (Word Error Rate)
    - CER  (Character Error Rate)
    - WER-norm (WER after unicode/punctuation/space normalization)
    - CER-norm (CER after normalization)
"""

import argparse
import json
import os
import re
import sys
import time
import unicodedata
from collections import defaultdict
from pathlib import Path

import numpy as np
import torch
import torchaudio
from datasets import load_dataset as hf_load_dataset
from jiwer import cer, wer
from transformers import AutoModelForImageTextToText, AutoProcessor

# Language code → display name mapping
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",
}

# ── Text normalization (standard ASR eval practice) ──────────────────────

def normalize_text(text: str) -> str:
    """Normalize text for fair WER/CER comparison.
    Unicode NFC → lowercase → remove punctuation → collapse whitespace → strip.
    """
    text = unicodedata.normalize("NFC", text)
    text = text.lower()
    text = re.sub(r"[^\w\s]", "", text, flags=re.UNICODE)
    text = re.sub(r"\s+", " ", text).strip()
    return text

def load_audio_array(sample, target_sr=16000):
    """Extract audio as numpy float32 array at target sample rate."""
    audio = sample["audio"]
    wav = np.array(audio["array"], dtype=np.float32)
    sr = audio["sampling_rate"]

    if sr != target_sr:
        wav_tensor = torch.from_numpy(wav).unsqueeze(0)
        wav_tensor = torchaudio.functional.resample(wav_tensor, sr, target_sr)
        wav = wav_tensor.squeeze(0).numpy()

    return wav

def run_inference(model, processor, dataset, batch_size: int = 16, max_new_tokens: int = 256):
    """Run batched inference on HF dataset and return list of result dicts."""
    results = []
    total = len(dataset)

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

        # Prepare batch: audio arrays + prompts
        wavs = []
        prompts = []
        batch_meta = []
        for sample in batch:
            lang_name = LANG_CODE_TO_NAME.get(sample["lang_code"], sample["language"].capitalize())
            user_content = f"<audio_soft_token>\nTranscribe the audio.\nLanguage: {lang_name}"
            messages = [{"role": "user", "content": user_content}]

            wavs.append(load_audio_array(sample))
            prompts.append(processor.apply_chat_template(
                messages, add_generation_prompt=True, tokenize=False
            ))
            batch_meta.append({
                "id": sample["id"],
                "language": sample["language"],
                "lang_code": sample["lang_code"],
                "source": sample["source"],
                "duration": sample["duration"],
                "reference": sample["reference"],
            })

        try:
            # Process batch
            inputs = processor(
                text=prompts, audio=wavs, sampling_rate=16000,
                return_tensors="pt", padding=True,
            )
            inputs = {k: v.to(model.device) if isinstance(v, torch.Tensor) else v
                      for k, v in inputs.items()}
            input_len = inputs["input_ids"].shape[1]

            # Generate
            with torch.no_grad():
                output_ids = model.generate(
                    **inputs, max_new_tokens=max_new_tokens, do_sample=False,
                )

            # Decode each sample in the batch
            for i in range(len(batch_meta)):
                gen_ids = output_ids[i][input_len:]
                hypothesis = processor.tokenizer.decode(gen_ids, skip_special_tokens=True).strip()
                results.append({**batch_meta[i], "hypothesis": hypothesis})

        except Exception as e:
            print(f"  Batch {start}-{end} FAILED ({e}), falling back to sequential...", flush=True)
            for i, sample in enumerate(batch):
                try:
                    inp = processor(
                        text=[prompts[i]], audio=[wavs[i]], sampling_rate=16000,
                        return_tensors="pt",
                    )
                    inp = {k: v.to(model.device) if isinstance(v, torch.Tensor) else v
                           for k, v in inp.items()}
                    with torch.no_grad():
                        out = model.generate(**inp, max_new_tokens=max_new_tokens, do_sample=False)
                    gen_ids = out[0][inp["input_ids"].shape[1]:]
                    hypothesis = processor.tokenizer.decode(gen_ids, skip_special_tokens=True).strip()
                except Exception as e2:
                    print(f"  Sample {start+i} ({batch_meta[i]['id']}) FAILED: {e2}")
                    hypothesis = ""
                results.append({**batch_meta[i], "hypothesis": hypothesis})

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

    return results

def compute_metrics(results: list) -> dict:
    """Compute WER, CER, normalized-WER/CER per language and overall."""
    by_lang = defaultdict(lambda: {"refs": [], "hyps": [], "refs_norm": [], "hyps_norm": []})

    for r in results:
        lang = r["language"]
        ref = r["reference"]
        hyp = r.get("hypothesis", "")

        ref_norm = normalize_text(ref)
        hyp_norm = normalize_text(hyp)

        if not ref_norm:
            continue

        by_lang[lang]["refs"].append(ref)
        by_lang[lang]["hyps"].append(hyp)
        by_lang[lang]["refs_norm"].append(ref_norm)
        by_lang[lang]["hyps_norm"].append(hyp_norm)

    metrics = {}
    all_refs, all_hyps = [], []
    all_refs_norm, all_hyps_norm = [], []

    for lang in sorted(by_lang.keys()):
        d = by_lang[lang]
        n = len(d["refs"])
        lang_wer = wer(d["refs"], d["hyps"])
        lang_cer = cer(d["refs"], d["hyps"])
        lang_wer_norm = wer(d["refs_norm"], d["hyps_norm"])
        lang_cer_norm = cer(d["refs_norm"], d["hyps_norm"])

        empty_count = sum(1 for h in d["hyps"] if not h.strip())

        metrics[lang] = {
            "n_samples": n,
            "wer": round(lang_wer * 100, 2),
            "cer": round(lang_cer * 100, 2),
            "wer_normalized": round(lang_wer_norm * 100, 2),
            "cer_normalized": round(lang_cer_norm * 100, 2),
            "empty_hypotheses": empty_count,
        }

        all_refs.extend(d["refs"])
        all_hyps.extend(d["hyps"])
        all_refs_norm.extend(d["refs_norm"])
        all_hyps_norm.extend(d["hyps_norm"])

    # Overall (corpus-level)
    metrics["__overall__"] = {
        "n_samples": len(all_refs),
        "wer": round(wer(all_refs, all_hyps) * 100, 2),
        "cer": round(cer(all_refs, all_hyps) * 100, 2),
        "wer_normalized": round(wer(all_refs_norm, all_hyps_norm) * 100, 2),
        "cer_normalized": round(cer(all_refs_norm, all_hyps_norm) * 100, 2),
    }

    # Macro average across languages
    lang_keys = [k for k in metrics if not k.startswith("__")]
    metrics["__macro_avg__"] = {
        "n_languages": len(lang_keys),
        "wer": round(np.mean([metrics[k]["wer"] for k in lang_keys]), 2),
        "cer": round(np.mean([metrics[k]["cer"] for k in lang_keys]), 2),
        "wer_normalized": round(np.mean([metrics[k]["wer_normalized"] for k in lang_keys]), 2),
        "cer_normalized": round(np.mean([metrics[k]["cer_normalized"] for k in lang_keys]), 2),
    }

    return metrics

def print_metrics_table(metrics: dict, checkpoint_name: str):
    """Pretty print metrics as a table."""
    print(f"\n{'='*85}")
    print(f"  BENCHMARK RESULTS: {checkpoint_name}")
    print(f"{'='*85}")
    print(f"{'Language':<14} {'N':>5} {'WER%':>8} {'CER%':>8} {'WER-n%':>8} {'CER-n%':>8}")
    print(f"{'-'*14} {'-'*5} {'-'*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']:>8.2f} {m['cer']:>8.2f} "
              f"{m['wer_normalized']:>8.2f} {m['cer_normalized']:>8.2f}")

    print(f"{'-'*14} {'-'*5} {'-'*8} {'-'*8} {'-'*8} {'-'*8}")

    ov = metrics["__overall__"]
    print(f"{'OVERALL':<14} {ov['n_samples']:>5} {ov['wer']:>8.2f} {ov['cer']:>8.2f} "
          f"{ov['wer_normalized']:>8.2f} {ov['cer_normalized']:>8.2f}")

    ma = metrics["__macro_avg__"]
    print(f"{'MACRO-AVG':<14} {ma['n_languages']:>5} {ma['wer']:>8.2f} {ma['cer']:>8.2f} "
          f"{ma['wer_normalized']:>8.2f} {ma['cer_normalized']:>8.2f}")
    print(f"{'='*85}\n")

def save_sample_analysis(results: list, output_dir: Path, n_per_lang: int = 5):
    """Save a few sample predictions per language for manual inspection."""
    by_lang = defaultdict(list)
    for r in results:
        by_lang[r["language"]].append(r)

    samples = []
    for lang in sorted(by_lang.keys()):
        lang_results = by_lang[lang]
        for r in lang_results[:n_per_lang]:
            samples.append({
                "id": r["id"],
                "language": r["language"],
                "reference": r["reference"],
                "hypothesis": r["hypothesis"],
                "ref_normalized": normalize_text(r["reference"]),
                "hyp_normalized": normalize_text(r["hypothesis"]),
            })

    path = output_dir / "sample_analysis.json"
    with open(path, "w", encoding="utf-8") as f:
        json.dump(samples, f, ensure_ascii=False, indent=2)
    print(f"Sample analysis saved: {path}")

def main():
    parser = argparse.ArgumentParser(description="Benchmark Gemma3n-E2B-ASR on Indic ASR dataset")
    parser.add_argument("--checkpoint", required=True, help="Path to model checkpoint directory")
    parser.add_argument("--checkpoint-name", default=None,
                        help="Name for this checkpoint in output files (default: dirname)")
    parser.add_argument("--hf-dataset", default="BayAreaBoys/indic-asr-benchmark-6k",
                        help="HuggingFace dataset name")
    parser.add_argument("--hf-token", default=None, help="HuggingFace token (or set HF_TOKEN env)")
    parser.add_argument("--cache-dir", default="/home/ubuntu/training/datasets/indic-asr-benchmark-6k")
    parser.add_argument("--output-dir", default="/home/ubuntu/training/benchmark_results_gemma3n")
    parser.add_argument("--batch-size", type=int, default=16,
                        help="Samples per batch for offline batched inference")
    parser.add_argument("--max-new-tokens", type=int, default=256)
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument("--languages", nargs="+", default=None,
                        help="Filter to specific language codes (e.g. hi en ta)")
    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 and processor
    print(f"Loading checkpoint: {args.checkpoint}")
    t0 = time.time()

    # Load tokenizer first, patch missing audio_token_id (version mismatch workaround)
    from transformers import AutoTokenizer, AutoConfig
    config = AutoConfig.from_pretrained(args.checkpoint, trust_remote_code=True)
    tokenizer = AutoTokenizer.from_pretrained(args.checkpoint, trust_remote_code=True, padding_side="left")

    # Gemma3nProcessor expects special token attrs on the tokenizer (added in transformers 5.x)
    # Patch them from config for compatibility with transformers 4.x
    _token_patches = {
        "audio_token_id": getattr(config, "audio_token_id", 262273),
        "audio_token": "<audio_soft_token>",
        "boa_token_id": getattr(config, "boa_token_id", 256000),
        "boa_token": "<start_of_audio>",
        "eoa_token_id": getattr(config, "eoa_token_id", 262272),
        "eoa_token": "<end_of_audio>",
        "boi_token_id": getattr(config, "boi_token_id", 255999),
        "boi_token": "<start_of_image>",
        "eoi_token_id": getattr(config, "eoi_token_id", 262144),
        "eoi_token": "<end_of_image>",
        "image_token_id": getattr(config, "image_token_id", 262145),
        "image_token": "<image_soft_token>",
    }
    for attr, val in _token_patches.items():
        if not hasattr(tokenizer, attr):
            setattr(tokenizer, attr, val)

    processor = AutoProcessor.from_pretrained(
        args.checkpoint, trust_remote_code=True, tokenizer=tokenizer
    )
    model = AutoModelForImageTextToText.from_pretrained(
        args.checkpoint,
        torch_dtype=torch.bfloat16,
        device_map=args.device,
        trust_remote_code=True,
        attn_implementation="sdpa",
    )
    model.eval()
    # Enable KV cache for faster autoregressive generation (training disables it)
    model.config.use_cache = True
    if hasattr(model.config, "text_config"):
        model.config.text_config.use_cache = True

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

    # Load dataset
    hf_token = args.hf_token or 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"
    )

    # Filter languages if specified
    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, dataset, args.batch_size, args.max_new_tokens)
    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 metrics
    metrics = compute_metrics(results)
    metrics["__meta__"] = {
        "checkpoint": args.checkpoint,
        "checkpoint_name": ckpt_name,
        "model_type": "gemma3n-E2B-asr",
        "dataset": args.hf_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),
    }

    metrics_file = output_dir / "metrics.json"
    with open(metrics_file, "w") as f:
        json.dump(metrics, f, indent=2, ensure_ascii=False)
    print(f"Metrics saved: {metrics_file}")

    # Print table
    print_metrics_table(metrics, ckpt_name)

    # Save sample analysis
    save_sample_analysis(results, output_dir)

if __name__ == "__main__":
    main()


[Request interrupted by user]


remember we are benchmarking qwen3-asr, not gemma.


Let me check the qwen3 benchmark script instead.


Bash

Find qwen3 benchmark script


1	#!/usr/bin/env python3
2	"""
3	Benchmark Qwen3-ASR checkpoints on indic-asr-benchmark-6k dataset.
4
5	Usage:
6	    python benchmark_qwen3_asr.py \
7	        --checkpoint /home/ubuntu/training/checkpoints/qwen3-asr-ckpt-24000 \
8	        --checkpoint-name ckpt-24000 \
9	        --output-dir /home/ubuntu/training/benchmark_results \
10	        --batch-size 8
11
12	Metrics computed per language and overall:
13	    - WER  (Word Error Rate)
14	    - CER  (Character Error Rate)
15	    - WER-norm (WER after unicode/punctuation/space normalization)
16	    - CER-norm (CER after normalization)
17	"""
18
19	import argparse
20	import json
21	import os
22	import re
23	import sys
24	import time
25	import unicodedata
26	from collections import defaultdict
27	from pathlib import Path
28
29	import numpy as np
30	import torch
31	from datasets import load_dataset as hf_load_dataset
32	from jiwer import cer, wer
33
34	# Add the repo to path for local imports
35	sys.path.insert(0, "/home/ubuntu/training/qwen3-asr-1.7b-phase2-sft")
36
37	# Language code → Qwen3-ASR language name mapping
38	LANG_CODE_TO_NAME = {
39	    "as": "Assamese", "bn": "Bengali", "en": "English", "gu": "Gujarati",
40	    "hi": "Hindi", "kn": "Kannada", "ml": "Malayalam", "mr": "Marathi",
41	    "or": "Odia", "pa": "Punjabi", "ta": "Tamil", "te": "Telugu",
42	}
43
44	# Languages the model supports (base + Indic from Phase 2 fine-tuning)
45	_MODEL_SUPPORTED_LANGS = {
46	    "Chinese", "English", "Cantonese", "Arabic", "German", "French", "Spanish",
47	    "Portuguese", "Indonesian", "Italian", "Korean", "Russian", "Thai", "Vietnamese",
48	    "Japanese", "Turkish", "Hindi", "Malay", "Dutch", "Swedish", "Danish", "Finnish",
49	    "Polish", "Czech", "Filipino", "Persian", "Greek", "Romanian", "Hungarian", "Macedonian",
50	    # Indic languages (Phase 2 fine-tuning)
51	    "Assamese", "Bengali", "Gujarati", "Kannada", "Malayalam",
52	    "Marathi", "Odia", "Punjabi", "Tamil", "Telugu",
53	}
54
55
56	# ── Text normalization (standard ASR eval practice) ──────────────────────
57
58	def normalize_text(text: str) -> str:
59	    """Normalize text for fair WER/CER comparison.
60	    Unicode NFC → lowercase → remove punctuation → collapse whitespace → strip.
61	    """
62	    text = unicodedata.normalize("NFC", text)
63	    text = text.lower()
64	    text = re.sub(r"[^\w\s]", "", text, flags=re.UNICODE)
65	    text = re.sub(r"\s+", " ", text).strip()
66	    return text
67
68
69	def run_inference(model, dataset, batch_size: int, force_language: bool = True):
70	    """Run batched inference on HF dataset and return list of result dicts."""
71	    results = []
72	    total = len(dataset)
73
74	    for start in range(0, total, batch_size):
75	        end = min(start + batch_size, total)
76	        batch = dataset.select(range(start, end))
77
78	        audio_inputs = []
79	        for sample in batch:
80	            audio = sample["audio"]
81	            wav = np.array(audio["array"], dtype=np.float32)
82	            sr = audio["sampling_rate"]
83	            audio_inputs.append((wav, sr))
84
85	        # Force language if supported
86	        if force_language:
87	            forced_langs = []
88	            for sample in batch:
89	                lang_name = LANG_CODE_TO_NAME.get(sample["lang_code"])
90	                forced_langs.append(lang_name)
91	        else:
92	            forced_langs = None
93
94	        try:
95	            transcriptions = model.transcribe(
96	                audio=audio_inputs,
97	                language=forced_langs,
98	                return_time_stamps=False,
99	            )
100	        except Exception as e:
101	            print(f"  Batch {start}-{end} failed ({e}), retrying without forced language...")
102	            try:
103	                transcriptions = model.transcribe(
104	                    audio=audio_inputs, language=None, return_time_stamps=False,
105	                )
106	            except Exception as e2:
107	                print(f"  Batch {start}-{end} FAILED: {e2}, trying one-by-one...")
108	                transcriptions = []
109	                for ai in audio_inputs:
110	                    try:
111	                        r = model.transcribe(audio=ai, language=None, return_time_stamps=False)
112	                        transcriptions.extend(r)
113	                    except Exception as e3:
114	                        from qwen_asr.inference.qwen3_asr import ASRTranscription
115	                        transcriptions.append(ASRTranscription(language="", text=""))
116
117	        for i, t in enumerate(transcriptions):
118	            sample = batch[i]
119	            results.append({
120	                "id": sample["id"],
121	                "language": sample["language"],
122	                "lang_code": sample["lang_code"],
123	                "source": sample["source"],
124	                "duration": sample["duration"],
125	                "reference": sample["reference"],
126	                "hypothesis": t.text,
127	                "detected_language": t.language,
128	            })
129
130	        elapsed_pct = end / total * 100
131	        print(f"  [{end}/{total}] ({elapsed_pct:.0f}%) processed", flush=True)
132
133	    return results
134
135
136	def compute_metrics(results: list) -> dict:
137	    """Compute WER, CER, normalized-WER/CER per language and overall."""
138	    by_lang = defaultdict(lambda: {"refs": [], "hyps": [], "refs_norm": [], "hyps_norm": []})
139
140	    for r in results:
141	        lang = r["language"]
142	        ref = r["reference"]
143	        hyp = r.get("hypothesis", "")
144
145	        ref_norm = normalize_text(ref)
146	        hyp_norm = normalize_text(hyp)
147
148	        if not ref_norm:
149	            continue
150
151	        by_lang[lang]["refs"].append(ref)
152	        by_lang[lang]["hyps"].append(hyp)
153	        by_lang[lang]["refs_norm"].append(ref_norm)
154	        by_lang[lang]["hyps_norm"].append(hyp_norm)
155
156	    metrics = {}
157	    all_refs, all_hyps = [], []
158	    all_refs_norm, all_hyps_norm = [], []
159
160	    for lang in sorted(by_lang.keys()):
161	        d = by_lang[lang]
162	        n = len(d["refs"])
163	        lang_wer = wer(d["refs"], d["hyps"])
164	        lang_cer = cer(d["refs"], d["hyps"])
165	        lang_wer_norm = wer(d["refs_norm"], d["hyps_norm"])
166	        lang_cer_norm = cer(d["refs_norm"], d["hyps_norm"])
167
168	        empty_count = sum(1 for h in d["hyps"] if not h.strip())
169
170	        metrics[lang] = {
171	            "n_samples": n,
172	            "wer": round(lang_wer * 100, 2),
173	            "cer": round(lang_cer * 100, 2),
174	            "wer_normalized": round(lang_wer_norm * 100, 2),
175	            "cer_normalized": round(lang_cer_norm * 100, 2),
176	            "empty_hypotheses": empty_count,
177	        }
178
179	        all_refs.extend(d["refs"])
180	        all_hyps.extend(d["hyps"])
181	        all_refs_norm.extend(d["refs_norm"])
182	        all_hyps_norm.extend(d["hyps_norm"])
183
184	    # Overall (corpus-level)
185	    metrics["__overall__"] = {
186	        "n_samples": len(all_refs),
187	        "wer": round(wer(all_refs, all_hyps) * 100, 2),
188	        "cer": round(cer(all_refs, all_hyps) * 100, 2),
189	        "wer_normalized": round(wer(all_refs_norm, all_hyps_norm) * 100, 2),
190	        "cer_normalized": round(cer(all_refs_norm, all_hyps_norm) * 100, 2),
191	    }
192
193	    # Macro average across languages
194	    lang_keys = [k for k in metrics if not k.startswith("__")]
195	    metrics["__macro_avg__"] = {
196	        "n_languages": len(lang_keys),
197	        "wer": round(np.mean([metrics[k]["wer"] for k in lang_keys]), 2),
198	        "cer": round(np.mean([metrics[k]["cer"] for k in lang_keys]), 2),
199	        "wer_normalized": round(np.mean([metrics[k]["wer_normalized"] for k in lang_keys]), 2),
200	        "cer_normalized": round(np.mean([metrics[k]["cer_normalized"] for k in lang_keys]), 2),
201	    }
202
203	    return metrics
204
205
206	def print_metrics_table(metrics: dict, checkpoint_name: str):
207	    """Pretty print metrics as a table."""
208	    print(f"\n{'='*85}")
209	    print(f"  BENCHMARK RESULTS: {checkpoint_name}")
210	    print(f"{'='*85}")
211	    print(f"{'Language':<14} {'N':>5} {'WER%':>8} {'CER%':>8} {'WER-n%':>8} {'CER-n%':>8}")
212	    print(f"{'-'*14} {'-'*5} {'-'*8} {'-'*8} {'-'*8} {'-'*8}")
213
214	    for lang in sorted(k for k in metrics if not k.startswith("__")):
215	        m = metrics[lang]
216	        print(f"{lang:<14} {m['n_samples']:>5} {m['wer']:>8.2f} {m['cer']:>8.2f} "
217	              f"{m['wer_normalized']:>8.2f} {m['cer_normalized']:>8.2f}")
218
219	    print(f"{'-'*14} {'-'*5} {'-'*8} {'-'*8} {'-'*8} {'-'*8}")
220
221	    ov = metrics["__overall__"]
222	    print(f"{'OVERALL':<14} {ov['n_samples']:>5} {ov['wer']:>8.2f} {ov['cer']:>8.2f} "
223	          f"{ov['wer_normalized']:>8.2f} {ov['cer_normalized']:>8.2f}")
224
225	    ma = metrics["__macro_avg__"]
226	    print(f"{'MACRO-AVG':<14} {ma['n_languages']:>5} {ma['wer']:>8.2f} {ma['cer']:>8.2f} "
227	          f"{ma['wer_normalized']:>8.2f} {ma['cer_normalized']:>8.2f}")
228	    print(f"{'='*85}\n")
229
230
231	def save_sample_analysis(results: list, output_dir: Path, n_per_lang: int = 5):
232	    """Save a few sample predictions per language for manual inspection."""
233	    by_lang = defaultdict(list)
234	    for r in results:
235	        by_lang[r["language"]].append(r)
236
237	    samples = []
238	    for lang in sorted(by_lang.keys()):
239	        lang_results = by_lang[lang]
240	        # Pick first n_per_lang samples
241	        for r in lang_results[:n_per_lang]:
242	            samples.append({
243	                "id": r["id"],
244	                "language": r["language"],
245	                "reference": r["reference"],
246	                "hypothesis": r["hypothesis"],
247	                "detected_language": r.get("detected_language", ""),
248	                "ref_normalized": normalize_text(r["reference"]),
249	                "hyp_normalized": normalize_text(r["hypothesis"]),
250	            })
251
252	    path = output_dir / "sample_analysis.json"
253	    with open(path, "w", encoding="utf-8") as f:
254	        json.dump(samples, f, ensure_ascii=False, indent=2)
255	    print(f"Sample analysis saved: {path}")
256
257
258	def main():
259	    parser = argparse.ArgumentParser(description="Benchmark Qwen3-ASR on Indic ASR dataset")
260	    parser.add_argument("--checkpoint", required=True, help="Path to model checkpoint directory")
261	    parser.add_argument("--checkpoint-name", default=None,
262	                        help="Name for this checkpoint in output files (default: dirname)")
263	    parser.add_argument("--hf-dataset", default="BayAreaBoys/indic-asr-benchmark-6k",
264	                        help="HuggingFace dataset name")
265	    parser.add_argument("--hf-token", default=None, help="HuggingFace token (or set HF_TOKEN env)")
266	    parser.add_argument("--cache-dir", default="/home/ubuntu/training/datasets/indic-asr-benchmark-6k")
267	    parser.add_argument("--output-dir", default="/home/ubuntu/training/benchmark_results")
268	    parser.add_argument("--batch-size", type=int, default=8)
269	    parser.add_argument("--max-new-tokens", type=int, default=512)
270	    parser.add_argument("--device", default="cuda:0")
271	    parser.add_argument("--backend", choices=["transformers", "vllm"], default="vllm",
272	                        help="Inference backend (default: vllm)")
273	    parser.add_argument("--gpu-memory-utilization", type=float, default=0.85,
274	                        help="vLLM GPU memory utilization (default: 0.85)")
275	    parser.add_argument("--no-force-language", action="store_true",
276	                        help="Don't force language, let model auto-detect")
277	    parser.add_argument("--languages", nargs="+", default=None,
278	                        help="Filter to specific language codes (e.g. hi en ta)")
279	    args = parser.parse_args()
280
281	    ckpt_name = args.checkpoint_name or Path(args.checkpoint).name
282	    output_dir = Path(args.output_dir) / ckpt_name
283	    output_dir.mkdir(parents=True, exist_ok=True)
284
285	    # Load model
286	    from qwen_asr import Qwen3ASRModel
287
288	    print(f"Loading checkpoint: {args.checkpoint} (backend={args.backend})")
289	    t0 = time.time()
290	    if args.backend == "vllm":
291	        model = Qwen3ASRModel.LLM(
292	            model=args.checkpoint,
293	            gpu_memory_utilization=args.gpu_memory_utilization,
294	            max_inference_batch_size=args.batch_size if args.batch_size > 0 else -1,
295	            max_new_tokens=args.max_new_tokens,
296	        )
297	    else:
298	        model = Qwen3ASRModel.from_pretrained(
299	            args.checkpoint,
300	            dtype=torch.bfloat16,
301	            device_map=args.device,
302	            max_inference_batch_size=args.batch_size,
303	            max_new_tokens=args.max_new_tokens,
304	        )
305	    load_time = time.time() - t0
306	    print(f"Model loaded in {load_time:.1f}s")
307
308	    # Load dataset
309	    hf_token = args.hf_token or os.environ.get("HF_TOKEN")
310	    print(f"Loading dataset: {args.hf_dataset}")
311	    dataset = hf_load_dataset(
312	        args.hf_dataset, token=hf_token, cache_dir=args.cache_dir, split="train"
313	    )
314
315	    # Filter languages if specified
316	    if args.languages:
317	        dataset = dataset.filter(lambda x: x["lang_code"] in args.languages)
318	        print(f"Filtered to {len(dataset)} samples for languages: {args.languages}")
319
320	    print(f"Dataset: {len(dataset)} samples")
321
322	    # Run inference
323	    print(f"\nRunning inference (batch_size={args.batch_size}, "
324	          f"force_language={not args.no_force_language})...")
325	    t0 = time.time()
326	    results = run_inference(model, dataset, args.batch_size,
327	                            force_language=not args.no_force_language)
328	    inference_time = time.time() - t0
329
330	    total_audio_sec = sum(r["duration"] for r in results)
331	    rtf = inference_time / total_audio_sec if total_audio_sec > 0 else 0
332	    print(f"Inference done in {inference_time:.1f}s for {total_audio_sec:.0f}s audio (RTF={rtf:.4f})")
333
334	    # Save raw predictions
335	    preds_file = output_dir / "predictions.json"
336	    with open(preds_file, "w", encoding="utf-8") as f:
337	        json.dump(results, f, ensure_ascii=False, indent=2)
338	    print(f"Predictions saved: {preds_file}")
339
340	    # Compute metrics
341	    metrics = compute_metrics(results)
342	    metrics["__meta__"] = {
343	        "checkpoint": args.checkpoint,
344	        "checkpoint_name": ckpt_name,
345	        "backend": args.backend,
346	        "dataset": args.hf_dataset,
347	        "batch_size": args.batch_size,
348	        "force_language": not args.no_force_language,
349	        "inference_time_sec": round(inference_time, 2),
350	        "total_audio_sec": round(total_audio_sec, 2),
351	        "rtf": round(rtf, 4),
352	    }
353
354	    metrics_file = output_dir / "metrics.json"
355	    with open(metrics_file, "w") as f:
356	        json.dump(metrics, f, indent=2, ensure_ascii=False)
357	    print(f"Metrics saved: {metrics_file}")
358
359	    # Print table
360	    print_metrics_table(metrics, ckpt_name)
361
362	    # Save sample analysis
363	    save_sample_analysis(results, output_dir)
364
365
366	if __name__ == "__main__":
367	    main()
368

#!/usr/bin/env python3
"""
Benchmark Qwen3-ASR checkpoints on indic-asr-benchmark-6k dataset.

Usage:
    python benchmark_qwen3_asr.py \
        --checkpoint /home/ubuntu/training/checkpoints/qwen3-asr-ckpt-24000 \
        --checkpoint-name ckpt-24000 \
        --output-dir /home/ubuntu/training/benchmark_results \
        --batch-size 8

Metrics computed per language and overall:
    - WER  (Word Error Rate)
    - CER  (Character Error Rate)
    - WER-norm (WER after unicode/punctuation/space normalization)
    - CER-norm (CER after normalization)
"""

import argparse
import json
import os
import re
import sys
import time
import unicodedata
from collections import defaultdict
from pathlib import Path

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

# Add the repo to path for local imports
sys.path.insert(0, "/home/ubuntu/training/qwen3-asr-1.7b-phase2-sft")

# Language code → Qwen3-ASR language name mapping
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",
}

# Languages the model supports (base + Indic from Phase 2 fine-tuning)
_MODEL_SUPPORTED_LANGS = {
    "Chinese", "English", "Cantonese", "Arabic", "German", "French", "Spanish",
    "Portuguese", "Indonesian", "Italian", "Korean", "Russian", "Thai", "Vietnamese",
    "Japanese", "Turkish", "Hindi", "Malay", "Dutch", "Swedish", "Danish", "Finnish",
    "Polish", "Czech", "Filipino", "Persian", "Greek", "Romanian", "Hungarian", "Macedonian",
    # Indic languages (Phase 2 fine-tuning)
    "Assamese", "Bengali", "Gujarati", "Kannada", "Malayalam",
    "Marathi", "Odia", "Punjabi", "Tamil", "Telugu",
}

# ── Text normalization (standard ASR eval practice) ──────────────────────

def normalize_text(text: str) -> str:
    """Normalize text for fair WER/CER comparison.
    Unicode NFC → lowercase → remove punctuation → collapse whitespace → strip.
    """
    text = unicodedata.normalize("NFC", text)
    text = text.lower()
    text = re.sub(r"[^\w\s]", "", text, flags=re.UNICODE)
    text = re.sub(r"\s+", " ", text).strip()
    return text

def run_inference(model, dataset, batch_size: int, force_language: bool = True):
    """Run batched inference on HF dataset and return list of result dicts."""
    results = []
    total = len(dataset)

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

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

        # Force language if supported
        if force_language:
            forced_langs = []
            for sample in batch:
                lang_name = LANG_CODE_TO_NAME.get(sample["lang_code"])
                forced_langs.append(lang_name)
        else:
            forced_langs = None

        try:
            transcriptions = model.transcribe(
                audio=audio_inputs,
                language=forced_langs,
                return_time_stamps=False,
            )
        except Exception as e:
            print(f"  Batch {start}-{end} failed ({e}), retrying without forced language...")
            try:
                transcriptions = model.transcribe(
                    audio=audio_inputs, language=None, return_time_stamps=False,
                )
            except Exception as e2:
                print(f"  Batch {start}-{end} FAILED: {e2}, trying one-by-one...")
                transcriptions = []
                for ai in audio_inputs:
                    try:
                        r = model.transcribe(audio=ai, language=None, return_time_stamps=False)
                        transcriptions.extend(r)
                    except Exception as e3:
                        from qwen_asr.inference.qwen3_asr import ASRTranscription
                        transcriptions.append(ASRTranscription(language="", text=""))

        for i, t in enumerate(transcriptions):
            sample = batch[i]
            results.append({
                "id": sample["id"],
                "language": sample["language"],
                "lang_code": sample["lang_code"],
                "source": sample["source"],
                "duration": sample["duration"],
                "reference": sample["reference"],
                "hypothesis": t.text,
                "detected_language": t.language,
            })

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

    return results

def compute_metrics(results: list) -> dict:
    """Compute WER, CER, normalized-WER/CER per language and overall."""
    by_lang = defaultdict(lambda: {"refs": [], "hyps": [], "refs_norm": [], "hyps_norm": []})

    for r in results:
        lang = r["language"]
        ref = r["reference"]
        hyp = r.get("hypothesis", "")

        ref_norm = normalize_text(ref)
        hyp_norm = normalize_text(hyp)

        if not ref_norm:
            continue

        by_lang[lang]["refs"].append(ref)
        by_lang[lang]["hyps"].append(hyp)
        by_lang[lang]["refs_norm"].append(ref_norm)
        by_lang[lang]["hyps_norm"].append(hyp_norm)

    metrics = {}
    all_refs, all_hyps = [], []
    all_refs_norm, all_hyps_norm = [], []

    for lang in sorted(by_lang.keys()):
        d = by_lang[lang]
        n = len(d["refs"])
        lang_wer = wer(d["refs"], d["hyps"])
        lang_cer = cer(d["refs"], d["hyps"])
        lang_wer_norm = wer(d["refs_norm"], d["hyps_norm"])
        lang_cer_norm = cer(d["refs_norm"], d["hyps_norm"])

        empty_count = sum(1 for h in d["hyps"] if not h.strip())

        metrics[lang] = {
            "n_samples": n,
            "wer": round(lang_wer * 100, 2),
            "cer": round(lang_cer * 100, 2),
            "wer_normalized": round(lang_wer_norm * 100, 2),
            "cer_normalized": round(lang_cer_norm * 100, 2),
            "empty_hypotheses": empty_count,
        }

        all_refs.extend(d["refs"])
        all_hyps.extend(d["hyps"])
        all_refs_norm.extend(d["refs_norm"])
        all_hyps_norm.extend(d["hyps_norm"])

    # Overall (corpus-level)
    metrics["__overall__"] = {
        "n_samples": len(all_refs),
        "wer": round(wer(all_refs, all_hyps) * 100, 2),
        "cer": round(cer(all_refs, all_hyps) * 100, 2),
        "wer_normalized": round(wer(all_refs_norm, all_hyps_norm) * 100, 2),
        "cer_normalized": round(cer(all_refs_norm, all_hyps_norm) * 100, 2),
    }

    # Macro average across languages
    lang_keys = [k for k in metrics if not k.startswith("__")]
    metrics["__macro_avg__"] = {
        "n_languages": len(lang_keys),
        "wer": round(np.mean([metrics[k]["wer"] for k in lang_keys]), 2),
        "cer": round(np.mean([metrics[k]["cer"] for k in lang_keys]), 2),
        "wer_normalized": round(np.mean([metrics[k]["wer_normalized"] for k in lang_keys]), 2),
        "cer_normalized": round(np.mean([metrics[k]["cer_normalized"] for k in lang_keys]), 2),
    }

    return metrics

def print_metrics_table(metrics: dict, checkpoint_name: str):
    """Pretty print metrics as a table."""
    print(f"\n{'='*85}")
    print(f"  BENCHMARK RESULTS: {checkpoint_name}")
    print(f"{'='*85}")
    print(f"{'Language':<14} {'N':>5} {'WER%':>8} {'CER%':>8} {'WER-n%':>8} {'CER-n%':>8}")
    print(f"{'-'*14} {'-'*5} {'-'*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']:>8.2f} {m['cer']:>8.2f} "
              f"{m['wer_normalized']:>8.2f} {m['cer_normalized']:>8.2f}")

    print(f"{'-'*14} {'-'*5} {'-'*8} {'-'*8} {'-'*8} {'-'*8}")

    ov = metrics["__overall__"]
    print(f"{'OVERALL':<14} {ov['n_samples']:>5} {ov['wer']:>8.2f} {ov['cer']:>8.2f} "
          f"{ov['wer_normalized']:>8.2f} {ov['cer_normalized']:>8.2f}")

    ma = metrics["__macro_avg__"]
    print(f"{'MACRO-AVG':<14} {ma['n_languages']:>5} {ma['wer']:>8.2f} {ma['cer']:>8.2f} "
          f"{ma['wer_normalized']:>8.2f} {ma['cer_normalized']:>8.2f}")
    print(f"{'='*85}\n")

def save_sample_analysis(results: list, output_dir: Path, n_per_lang: int = 5):
    """Save a few sample predictions per language for manual inspection."""
    by_lang = defaultdict(list)
    for r in results:
        by_lang[r["language"]].append(r)

    samples = []
    for lang in sorted(by_lang.keys()):
        lang_results = by_lang[lang]
        # Pick first n_per_lang samples
        for r in lang_results[:n_per_lang]:
            samples.append({
                "id": r["id"],
                "language": r["language"],
                "reference": r["reference"],
                "hypothesis": r["hypothesis"],
                "detected_language": r.get("detected_language", ""),
                "ref_normalized": normalize_text(r["reference"]),
                "hyp_normalized": normalize_text(r["hypothesis"]),
            })

    path = output_dir / "sample_analysis.json"
    with open(path, "w", encoding="utf-8") as f:
        json.dump(samples, f, ensure_ascii=False, indent=2)
    print(f"Sample analysis saved: {path}")

def main():
    parser = argparse.ArgumentParser(description="Benchmark Qwen3-ASR on Indic ASR dataset")
    parser.add_argument("--checkpoint", required=True, help="Path to model checkpoint directory")
    parser.add_argument("--checkpoint-name", default=None,
                        help="Name for this checkpoint in output files (default: dirname)")
    parser.add_argument("--hf-dataset", default="BayAreaBoys/indic-asr-benchmark-6k",
                        help="HuggingFace dataset name")
    parser.add_argument("--hf-token", default=None, help="HuggingFace token (or set HF_TOKEN env)")
    parser.add_argument("--cache-dir", default="/home/ubuntu/training/datasets/indic-asr-benchmark-6k")
    parser.add_argument("--output-dir", default="/home/ubuntu/training/benchmark_results")
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--max-new-tokens", type=int, default=512)
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument("--backend", choices=["transformers", "vllm"], default="vllm",
                        help="Inference backend (default: vllm)")
    parser.add_argument("--gpu-memory-utilization", type=float, default=0.85,
                        help="vLLM GPU memory utilization (default: 0.85)")
    parser.add_argument("--no-force-language", action="store_true",
                        help="Don't force language, let model auto-detect")
    parser.add_argument("--languages", nargs="+", default=None,
                        help="Filter to specific language codes (e.g. hi en ta)")
    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 qwen_asr import Qwen3ASRModel

    print(f"Loading checkpoint: {args.checkpoint} (backend={args.backend})")
    t0 = time.time()
    if args.backend == "vllm":
        model = Qwen3ASRModel.LLM(
            model=args.checkpoint,
            gpu_memory_utilization=args.gpu_memory_utilization,
            max_inference_batch_size=args.batch_size if args.batch_size > 0 else -1,
            max_new_tokens=args.max_new_tokens,
        )
    else:
        model = Qwen3ASRModel.from_pretrained(
            args.checkpoint,
            dtype=torch.bfloat16,
            device_map=args.device,
            max_inference_batch_size=args.batch_size,
            max_new_tokens=args.max_new_tokens,
        )
    load_time = time.time() - t0
    print(f"Model loaded in {load_time:.1f}s")

    # Load dataset
    hf_token = args.hf_token or 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"
    )

    # Filter languages if specified
    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}, "
          f"force_language={not args.no_force_language})...")
    t0 = time.time()
    results = run_inference(model, dataset, args.batch_size,
                            force_language=not args.no_force_language)
    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 metrics
    metrics = compute_metrics(results)
    metrics["__meta__"] = {
        "checkpoint": args.checkpoint,
        "checkpoint_name": ckpt_name,
        "backend": args.backend,
        "dataset": args.hf_dataset,
        "batch_size": args.batch_size,
        "force_language": not args.no_force_language,
        "inference_time_sec": round(inference_time, 2),
        "total_audio_sec": round(total_audio_sec, 2),
        "rtf": round(rtf, 4),
    }

    metrics_file = output_dir / "metrics.json"
    with open(metrics_file, "w") as f:
        json.dump(metrics, f, indent=2, ensure_ascii=False)
    print(f"Metrics saved: {metrics_file}")

    # Print table
    print_metrics_table(metrics, ckpt_name)

    # Save sample analysis
    save_sample_analysis(results, output_dir)

if __name__ == "__main__":
    main()


Wait — this is the **qwen3** benchmark script, not gemma3n. You said the `--backend vllm` route is broken. Let me check what `Qwen3ASRModel.LLM` does and find the actual error.


python3 -c "
import sys
sys.path.insert(0, '/home/ubuntu/training/qwen3-asr-1.7b-phase2-sft')
from qwen_asr import Qwen3ASRModel
# Check what LLM method expects
import inspect
print(inspect.signature(Qwen3ASRModel.LLM))
print()
src = inspect.getsource(Qwen3ASRModel.LLM)
print(src[:3000])
" 2>&1

Check Qwen3ASRModel.LLM signature and source


Skipping import of cpp extensions due to incompatible torch version 2.9.1+cu128 for torchao version 0.16.0             Please see https://github.com/pytorch/ao/issues/2919 for more info
WARNING 03-31 15:06:57 [registry.py:847] Model architecture Qwen3ASRForConditionalGeneration is already registered, and will be overwritten by the new model class <class 'qwen_asr.core.vllm_backend.qwen3_asr.Qwen3ASRForConditionalGeneration'>.
(model: str, forced_aligner: Optional[str] = None, forced_aligner_kwargs: Optional[Dict[str, Any]] = None, max_inference_batch_size: int = -1, max_new_tokens: Optional[int] = 4096, **kwargs) -> 'Qwen3ASRModel'

    @classmethod
    def LLM(
        cls,
        model: str,
        forced_aligner: Optional[str] = None,
        forced_aligner_kwargs: Optional[Dict[str, Any]] = None,
        max_inference_batch_size: int = -1,
        max_new_tokens: Optional[int] = 4096,
        **kwargs,
    ) -> "Qwen3ASRModel":
        """
        Initialize using vLLM backend.

        Import is isolated to keep vLLM optional.

        Args:
            model:
                Model path/repo for vLLM.
            forced_aligner:
                Optional forced aligner model path/repo id.
            forced_aligner_kwargs:
                Optional kwargs forwarded to Qwen3ForcedAligner.from_pretrained(...).
            max_inference_batch_size:
                Batch size limit for inference. -1 means no chunking. Small values can avoid OOM.
            max_new_tokens:
                Maximum number of tokens to generate.
            **kwargs:
                Forwarded to vllm.LLM(...).

        Returns:
            Qwen3ASRModel

        Raises:
            ImportError: If vLLM is not installed.
        """
        try:
            from vllm import LLM as vLLM
            from vllm import SamplingParams
        except Exception as e:
            raise ImportError(
                "vLLM is not available. Install with: pip install qwen-asr[vllm]"
            ) from e

        llm = vLLM(model=model, **kwargs)

        processor = Qwen3ASRProcessor.from_pretrained(model, fix_mistral_regex=True)
        sampling_params = SamplingParams(**({"temperature": 0.0, "max_tokens": max_new_tokens}))

        forced_aligner_model = None
        if forced_aligner is not None:
            forced_aligner_model = Qwen3ForcedAligner.from_pretrained(
                forced_aligner, **(forced_aligner_kwargs or {})
            )

        return cls(
            backend="vllm",
            model=llm,
            processor=processor,
            sampling_params=sampling_params,
            forced_aligner=forced_aligner_model,
            max_inference_batch_size=max_inference_batch_size,
            max_new_tokens=None,
        )


python3 benchmark_qwen3_asr.py \
    --checkpoint /home/ubuntu/training/checkpoints/qwen3-asr-ckpt-400000 \
    --checkpoint-name test-vllm \
    --output-dir /tmp/test-vllm \
    --backend vllm \
    --batch-size 32 \
    --languages hi \
    2>&1 | head -60

Test vLLM backend with Hindi only to see the error


Skipping import of cpp extensions due to incompatible torch version 2.9.1+cu128 for torchao version 0.16.0             Please see https://github.com/pytorch/ao/issues/2919 for more info
WARNING 03-31 15:07:19 [registry.py:847] Model architecture Qwen3ASRForConditionalGeneration is already registered, and will be overwritten by the new model class <class 'qwen_asr.core.vllm_backend.qwen3_asr.Qwen3ASRForConditionalGeneration'>.
Loading checkpoint: /home/ubuntu/training/checkpoints/qwen3-asr-ckpt-400000 (backend=vllm)
INFO 03-31 15:07:20 [utils.py:223] non-default args: {'gpu_memory_utilization': 0.85, 'disable_log_stats': True, 'model': '/home/ubuntu/training/checkpoints/qwen3-asr-ckpt-400000'}
Unrecognized keys in `rope_scaling` for 'rope_type'='default': {'mrope_section', 'interleaved', 'mrope_interleaved'}
INFO 03-31 15:07:20 [qwen3_asr.py:414] thinker_config is None. Initializing thinker model with default values
INFO 03-31 15:07:20 [model.py:529] Resolved architecture: Qwen3ASRForConditionalGeneration
ERROR 03-31 15:07:20 [repo_utils.py:47] Error retrieving safetensors: Repo id must be in the form 'repo_name' or 'namespace/repo_name': '/home/ubuntu/training/checkpoints/qwen3-asr-ckpt-400000'. Use `repo_type` argument if needed., retrying 1 of 2
ERROR 03-31 15:07:22 [repo_utils.py:45] Error retrieving safetensors: Repo id must be in the form 'repo_name' or 'namespace/repo_name': '/home/ubuntu/training/checkpoints/qwen3-asr-ckpt-400000'. Use `repo_type` argument if needed.
INFO 03-31 15:07:22 [model.py:1871] Downcasting torch.float32 to torch.bfloat16.
INFO 03-31 15:07:22 [model.py:1549] Using max model len 65536
INFO 03-31 15:07:22 [scheduler.py:224] Chunked prefill is enabled with max_num_batched_tokens=8192.
INFO 03-31 15:07:22 [vllm.py:689] Asynchronous scheduling is enabled.
The following generation flags are not valid and may be ignored: ['temperature']. Set `TRANSFORMERS_VERBOSITY=info` for more details.
The tokenizer you are loading from '/home/ubuntu/training/checkpoints/qwen3-asr-ckpt-400000' with an incorrect regex pattern: https://huggingface.co/mistralai/Mistral-Small-3.1-24B-Instruct-2503/discussions/84#69121093e8b480e709447d5e. This will lead to incorrect tokenization. You should set the `fix_mistral_regex=True` flag when loading this tokenizer to fix this issue.
The tokenizer you are loading from '/home/ubuntu/training/checkpoints/qwen3-asr-ckpt-400000' with an incorrect regex pattern: https://huggingface.co/mistralai/Mistral-Small-3.1-24B-Instruct-2503/discussions/84#69121093e8b480e709447d5e. This will lead to incorrect tokenization. You should set the `fix_mistral_regex=True` flag when loading this tokenizer to fix this issue.
The tokenizer you are loading from '/home/ubuntu/training/checkpoints/qwen3-asr-ckpt-400000' with an incorrect regex pattern: https://huggingface.co/mistralai/Mistral-Small-3.1-24B-Instruct-2503/discussions/84#69121093e8b480e709447d5e. This will lead to incorrect tokenization. You should set the `fix_mistral_regex=True` flag when loading this tokenizer to fix this issue.
The tokenizer you are loading from '/home/ubuntu/training/checkpoints/qwen3-asr-ckpt-400000' with an incorrect regex pattern: https://huggingface.co/mistralai/Mistral-Small-3.1-24B-Instruct-2503/discussions/84#69121093e8b480e709447d5e. This will lead to incorrect tokenization. You should set the `fix_mistral_regex=True` flag when loading this tokenizer to fix this issue.
(EngineCore_DP0 pid=3279062) INFO 03-31 15:07:27 [core.py:97] Initializing a V1 LLM engine (v0.16.0) with config: model='/home/ubuntu/training/checkpoints/qwen3-asr-ckpt-400000', speculative_config=None, tokenizer='/home/ubuntu/training/checkpoints/qwen3-asr-ckpt-400000', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.bfloat16, max_seq_len=65536, download_dir=None, load_format=auto, tensor_parallel_size=1, pipeline_parallel_size=1, data_parallel_size=1, disable_custom_all_reduce=False, quantization=None, enforce_eager=False, enable_return_routed_experts=False, kv_cache_dtype=auto, device_config=cuda, structured_outputs_config=StructuredOutputsConfig(backend='auto', disable_fallback=False, disable_any_whitespace=False, disable_additional_properties=False, reasoning_parser='', reasoning_parser_plugin='', enable_in_reasoning=False), observability_config=ObservabilityConfig(show_hidden_metrics_for_version=None, otlp_traces_endpoint=None, collect_detailed_traces=None, kv_cache_metrics=False, kv_cache_metrics_sample=0.01, cudagraph_metrics=False, enable_layerwise_nvtx_tracing=False, enable_mfu_metrics=False, enable_mm_processor_stats=False, enable_logging_iteration_details=False), seed=0, served_model_name=/home/ubuntu/training/checkpoints/qwen3-asr-ckpt-400000, enable_prefix_caching=True, enable_chunked_prefill=True, pooler_config=None, compilation_config={'level': None, 'mode': <CompilationMode.VLLM_COMPILE: 3>, 'debug_dump_path': None, 'cache_dir': '', 'compile_cache_save_format': 'binary', 'backend': 'inductor', 'custom_ops': ['none'], 'splitting_ops': ['vllm::unified_attention', 'vllm::unified_attention_with_output', 'vllm::unified_mla_attention', 'vllm::unified_mla_attention_with_output', 'vllm::mamba_mixer2', 'vllm::mamba_mixer', 'vllm::short_conv', 'vllm::linear_attention', 'vllm::plamo2_mamba_mixer', 'vllm::gdn_attention_core', 'vllm::kda_attention', 'vllm::sparse_attn_indexer', 'vllm::rocm_aiter_sparse_attn_indexer', 'vllm::unified_kv_cache_update'], 'compile_mm_encoder': False, 'compile_sizes': [], 'compile_ranges_split_points': [8192], 'inductor_compile_config': {'enable_auto_functionalized_v2': False, 'combo_kernels': True, 'benchmark_combo_kernel': True}, 'inductor_passes': {}, 'cudagraph_mode': <CUDAGraphMode.FULL_AND_PIECEWISE: (2, 1)>, 'cudagraph_num_of_warmups': 1, 'cudagraph_capture_sizes': [1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256, 272, 288, 304, 320, 336, 352, 368, 384, 400, 416, 432, 448, 464, 480, 496, 512], 'cudagraph_copy_inputs': False, 'cudagraph_specialize_lora': True, 'use_inductor_graph_partition': False, 'pass_config': {'fuse_norm_quant': False, 'fuse_act_quant': False, 'fuse_attn_quant': False, 'eliminate_noops': True, 'enable_sp': False, 'fuse_gemm_comms': False, 'fuse_allreduce_rms': False, 'fuse_act_padding': False}, 'max_cudagraph_capture_size': 512, 'dynamic_shapes_config': {'type': <DynamicShapesType.BACKED: 'backed'>, 'evaluate_guards': False, 'assume_32_bit_indexing': False}, 'local_cache_dir': None, 'fast_moe_cold_start': True, 'static_all_moe_layers': []}
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006] EngineCore failed to start.
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006] Traceback (most recent call last):
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006]   File "/home/ubuntu/.local/lib/python3.10/site-packages/vllm/v1/engine/core.py", line 996, in run_engine_core
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006]     engine_core = EngineCoreProc(*args, engine_index=dp_rank, **kwargs)
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006]   File "/home/ubuntu/.local/lib/python3.10/site-packages/vllm/tracing/otel.py", line 178, in sync_wrapper
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006]     return func(*args, **kwargs)
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006]   File "/home/ubuntu/.local/lib/python3.10/site-packages/vllm/v1/engine/core.py", line 740, in __init__
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006]     super().__init__(
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006]   File "/home/ubuntu/.local/lib/python3.10/site-packages/vllm/v1/engine/core.py", line 106, in __init__
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006]     self.model_executor = executor_class(vllm_config)
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006]   File "/home/ubuntu/.local/lib/python3.10/site-packages/vllm/tracing/otel.py", line 178, in sync_wrapper
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006]     return func(*args, **kwargs)
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006]   File "/home/ubuntu/.local/lib/python3.10/site-packages/vllm/v1/executor/abstract.py", line 103, in __init__
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006]     self._init_executor()
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006]   File "/home/ubuntu/.local/lib/python3.10/site-packages/vllm/v1/executor/uniproc_executor.py", line 47, in _init_executor
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006]     self.driver_worker.init_device()
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006]   File "/home/ubuntu/.local/lib/python3.10/site-packages/vllm/v1/worker/worker_base.py", line 322, in init_device
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006]     self.worker.init_device()  # type: ignore
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006]   File "/home/ubuntu/.local/lib/python3.10/site-packages/vllm/tracing/otel.py", line 178, in sync_wrapper
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006]     return func(*args, **kwargs)
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006]   File "/home/ubuntu/.local/lib/python3.10/site-packages/vllm/v1/worker/gpu_worker.py", line 224, in init_device
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006]     current_platform.set_device(self.device)
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006]   File "/home/ubuntu/.local/lib/python3.10/site-packages/vllm/platforms/cuda.py", line 126, in set_device
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006]     torch.cuda.set_device(device)
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006]   File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/cuda/__init__.py", line 567, in set_device
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006]     torch._C._cuda_setDevice(device)
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006]   File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/cuda/__init__.py", line 398, in _lazy_init
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006]     raise RuntimeError(
(EngineCore_DP0 pid=3279062) ERROR 03-31 15:07:27 [core.py:1006] RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method
(EngineCore_DP0 pid=3279062) Process EngineCore_DP0:
(EngineCore_DP0 pid=3279062) Traceback (most recent call last):
(EngineCore_DP0 pid=3279062)   File "/usr/lib/python3.10/multiprocessing/process.py", line 314, in _bootstrap
(EngineCore_DP0 pid=3279062)     self.run()
(EngineCore_DP0 pid=3279062)   File "/usr/lib/python3.10/multiprocessing/process.py", line 108, in run
(EngineCore_DP0 pid=3279062)     self._target(*self._args, **self._kwargs)
(EngineCore_DP0 pid=3279062)   File "/home/ubuntu/.local/lib/python3.10/site-packages/vllm/v1/engine/core.py", line 1010, in run_engine_core
(EngineCore_DP0 pid=3279062)     raise e
(EngineCore_DP0 pid=3279062)   File "/home/ubuntu/.local/lib/python3.10/site-packages/vllm/v1/engine/core.py", line 996, in run_engine_core
(EngineCore_DP0 pid=3279062)     engine_core = EngineCoreProc(*args, engine_index=dp_rank, **kwargs)
(EngineCore_DP0 pid=3279062)   File "/home/ubuntu/.local/lib/python3.10/site-packages/vllm/tracing/otel.py", line 178, in sync_wrapper
(EngineCore_DP0 pid=3279062)     return func(*args, **kwargs)


Found it. The error is:

```
RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method
```

vLLM v1 engine uses multiprocessing but CUDA was already initialized in the parent process (probably by the `import torch` at the top of the script). The fix is to either:
1. Set `VLLM_WORKER_MULTIPROC_METHOD=spawn` env var, or
2. Avoid touching CUDA before vLLM init

Let me also check if there's a simpler vLLM config fix.


python3 -c "import vllm; print(vllm.__version__)" 2>&1

Check vLLM version


vLLM 0.16.0 uses V1 engine by default which spawns a subprocess. The fix is straightforward — set the multiprocessing start method to `spawn` before vLLM initializes. Let me also defer the `torch` import so CUDA isn't initialized early.


Edit

import argparse
import json
import os
import re
import sys
import time
import unicodedata
from collections import defaultdict
from pathlib import Path

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

# Add the repo to path for local imports
sys.path.insert(0, "/home/ubuntu/training/qwen3-asr-1.7b-phase2-sft")

import argparse
import json
import os
import re
import sys
import time
import unicodedata
from collections import defaultdict
from pathlib import Path

# vLLM v1 (0.16+) spawns a subprocess for the engine — CUDA must not be
# initialized in the parent before that fork.  Setting the env var up-front
# tells vLLM to use 'spawn' instead of 'fork', which avoids the
# "Cannot re-initialize CUDA in forked subprocess" crash.
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")

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

# Add the repo to path for local imports
sys.path.insert(0, "/home/ubuntu/training/qwen3-asr-1.7b-phase2-sft")


The file /home/ubuntu/training/benchmark_qwen3_asr.py has been updated successfully.

import argparse
import json
import os
import re
import sys
import time
import unicodedata
from collections import defaultdict
from pathlib import Path

# vLLM v1 (0.16+) spawns a subprocess for the engine — CUDA must not be
# initialized in the parent before that fork.  Setting the env var up-front
# tells vLLM to use 'spawn' instead of 'fork', which avoids the
# "Cannot re-initialize CUDA in forked subprocess" crash.
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")

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

# Add the repo to path for local imports
sys.path.insert(0, "/home/ubuntu/training/qwen3-asr-1.7b-phase2-sft")

#!/usr/bin/env python3
"""
Benchmark Qwen3-ASR checkpoints on indic-asr-benchmark-6k dataset.

Usage:
    python benchmark_qwen3_asr.py \
        --checkpoint /home/ubuntu/training/checkpoints/qwen3-asr-ckpt-24000 \
        --checkpoint-name ckpt-24000 \
        --output-dir /home/ubuntu/training/benchmark_results \
        --batch-size 8

Metrics computed per language and overall:
    - WER  (Word Error Rate)
    - CER  (Character Error Rate)
    - WER-norm (WER after unicode/punctuation/space normalization)
    - CER-norm (CER after normalization)
"""

import argparse
import json
import os
import re
import sys
import time
import unicodedata
from collections import defaultdict
from pathlib import Path

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

# Add the repo to path for local imports
sys.path.insert(0, "/home/ubuntu/training/qwen3-asr-1.7b-phase2-sft")

# Language code → Qwen3-ASR language name mapping
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",
}

# Languages the model supports (base + Indic from Phase 2 fine-tuning)
_MODEL_SUPPORTED_LANGS = {
    "Chinese", "English", "Cantonese", "Arabic", "German", "French", "Spanish",
    "Portuguese", "Indonesian", "Italian", "Korean", "Russian", "Thai", "Vietnamese",
    "Japanese", "Turkish", "Hindi", "Malay", "Dutch", "Swedish", "Danish", "Finnish",
    "Polish", "Czech", "Filipino", "Persian", "Greek", "Romanian", "Hungarian", "Macedonian",
    # Indic languages (Phase 2 fine-tuning)
    "Assamese", "Bengali", "Gujarati", "Kannada", "Malayalam",
    "Marathi", "Odia", "Punjabi", "Tamil", "Telugu",
}

# ── Text normalization (standard ASR eval practice) ──────────────────────

def normalize_text(text: str) -> str:
    """Normalize text for fair WER/CER comparison.
    Unicode NFC → lowercase → remove punctuation → collapse whitespace → strip.
    """
    text = unicodedata.normalize("NFC", text)
    text = text.lower()
    text = re.sub(r"[^\w\s]", "", text, flags=re.UNICODE)
    text = re.sub(r"\s+", " ", text).strip()
    return text

def run_inference(model, dataset, batch_size: int, force_language: bool = True):
    """Run batched inference on HF dataset and return list of result dicts."""
    results = []
    total = len(dataset)

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

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

        # Force language if supported
        if force_language:
            forced_langs = []
            for sample in batch:
                lang_name = LANG_CODE_TO_NAME.get(sample["lang_code"])
                forced_langs.append(lang_name)
        else:
            forced_langs = None

        try:
            transcriptions = model.transcribe(
                audio=audio_inputs,
                language=forced_langs,
                return_time_stamps=False,
            )
        except Exception as e:
            print(f"  Batch {start}-{end} failed ({e}), retrying without forced language...")
            try:
                transcriptions = model.transcribe(
                    audio=audio_inputs, language=None, return_time_stamps=False,
                )
            except Exception as e2:
                print(f"  Batch {start}-{end} FAILED: {e2}, trying one-by-one...")
                transcriptions = []
                for ai in audio_inputs:
                    try:
                        r = model.transcribe(audio=ai, language=None, return_time_stamps=False)
                        transcriptions.extend(r)
                    except Exception as e3:
                        from qwen_asr.inference.qwen3_asr import ASRTranscription
                        transcriptions.append(ASRTranscription(language="", text=""))

        for i, t in enumerate(transcriptions):
            sample = batch[i]
            results.append({
                "id": sample["id"],
                "language": sample["language"],
                "lang_code": sample["lang_code"],
                "source": sample["source"],
                "duration": sample["duration"],
                "reference": sample["reference"],
                "hypothesis": t.text,
                "detected_language": t.language,
            })

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

    return results

def compute_metrics(results: list) -> dict:
    """Compute WER, CER, normalized-WER/CER per language and overall."""
    by_lang = defaultdict(lambda: {"refs": [], "hyps": [], "refs_norm": [], "hyps_norm": []})

    for r in results:
        lang = r["language"]
        ref = r["reference"]
        hyp = r.get("hypothesis", "")

        ref_norm = normalize_text(ref)
        hyp_norm = normalize_text(hyp)

        if not ref_norm:
            continue

        by_lang[lang]["refs"].append(ref)
        by_lang[lang]["hyps"].append(hyp)
        by_lang[lang]["refs_norm"].append(ref_norm)
        by_lang[lang]["hyps_norm"].append(hyp_norm)

    metrics = {}
    all_refs, all_hyps = [], []
    all_refs_norm, all_hyps_norm = [], []

    for lang in sorted(by_lang.keys()):
        d = by_lang[lang]
        n = len(d["refs"])
        lang_wer = wer(d["refs"], d["hyps"])
        lang_cer = cer(d["refs"], d["hyps"])
        lang_wer_norm = wer(d["refs_norm"], d["hyps_norm"])
        lang_cer_norm = cer(d["refs_norm"], d["hyps_norm"])

        empty_count = sum(1 for h in d["hyps"] if not h.strip())

        metrics[lang] = {
            "n_samples": n,
            "wer": round(lang_wer * 100, 2),
            "cer": round(lang_cer * 100, 2),
            "wer_normalized": round(lang_wer_norm * 100, 2),
            "cer_normalized": round(lang_cer_norm * 100, 2),
            "empty_hypotheses": empty_count,
        }

        all_refs.extend(d["refs"])
        all_hyps.extend(d["hyps"])
        all_refs_norm.extend(d["refs_norm"])
        all_hyps_norm.extend(d["hyps_norm"])

    # Overall (corpus-level)
    metrics["__overall__"] = {
        "n_samples": len(all_refs),
        "wer": round(wer(all_refs, all_hyps) * 100, 2),
        "cer": round(cer(all_refs, all_hyps) * 100, 2),
        "wer_normalized": round(wer(all_refs_norm, all_hyps_norm) * 100, 2),
        "cer_normalized": round(cer(all_refs_norm, all_hyps_norm) * 100, 2),
    }

    # Macro average across languages
    lang_keys = [k for k in metrics if not k.startswith("__")]
    metrics["__macro_avg__"] = {
        "n_languages": len(lang_keys),
        "wer": round(np.mean([metrics[k]["wer"] for k in lang_keys]), 2),
        "cer": round(np.mean([metrics[k]["cer"] for k in lang_keys]), 2),
        "wer_normalized": round(np.mean([metrics[k]["wer_normalized"] for k in lang_keys]), 2),
        "cer_normalized": round(np.mean([metrics[k]["cer_normalized"] for k in lang_keys]), 2),
    }

    return metrics

def print_metrics_table(metrics: dict, checkpoint_name: str):
    """Pretty print metrics as a table."""
    print(f"\n{'='*85}")
    print(f"  BENCHMARK RESULTS: {checkpoint_name}")
    print(f"{'='*85}")
    print(f"{'Language':<14} {'N':>5} {'WER%':>8} {'CER%':>8} {'WER-n%':>8} {'CER-n%':>8}")
    print(f"{'-'*14} {'-'*5} {'-'*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']:>8.2f} {m['cer']:>8.2f} "
              f"{m['wer_normalized']:>8.2f} {m['cer_normalized']:>8.2f}")

    print(f"{'-'*14} {'-'*5} {'-'*8} {'-'*8} {'-'*8} {'-'*8}")

    ov = metrics["__overall__"]
    print(f"{'OVERALL':<14} {ov['n_samples']:>5} {ov['wer']:>8.2f} {ov['cer']:>8.2f} "
          f"{ov['wer_normalized']:>8.2f} {ov['cer_normalized']:>8.2f}")

    ma = metrics["__macro_avg__"]
    print(f"{'MACRO-AVG':<14} {ma['n_languages']:>5} {ma['wer']:>8.2f} {ma['cer']:>8.2f} "
          f"{ma['wer_normalized']:>8.2f} {ma['cer_normalized']:>8.2f}")
    print(f"{'='*85}\n")

def save_sample_analysis(results: list, output_dir: Path, n_per_lang: int = 5):
    """Save a few sample predictions per language for manual inspection."""
    by_lang = defaultdict(list)
    for r in results:
        by_lang[r["language"]].append(r)

    samples = []
    for lang in sorted(by_lang.keys()):
        lang_results = by_lang[lang]
        # Pick first n_per_lang samples
        for r in lang_results[:n_per_lang]:
            samples.append({
                "id": r["id"],
                "language": r["language"],
                "reference": r["reference"],
                "hypothesis": r["hypothesis"],
                "detected_language": r.get("detected_language", ""),
                "ref_normalized": normalize_text(r["reference"]),
                "hyp_normalized": normalize_text(r["hypothesis"]),
            })

    path = output_dir / "sample_analysis.json"
    with open(path, "w", encoding="utf-8") as f:
        json.dump(samples, f, ensure_ascii=False, indent=2)
    print(f"Sample analysis saved: {path}")

def main():
    parser = argparse.ArgumentParser(description="Benchmark Qwen3-ASR on Indic ASR dataset")
    parser.add_argument("--checkpoint", required=True, help="Path to model checkpoint directory")
    parser.add_argument("--checkpoint-name", default=None,
                        help="Name for this checkpoint in output files (default: dirname)")
    parser.add_argument("--hf-dataset", default="BayAreaBoys/indic-asr-benchmark-6k",
                        help="HuggingFace dataset name")
    parser.add_argument("--hf-token", default=None, help="HuggingFace token (or set HF_TOKEN env)")
    parser.add_argument("--cache-dir", default="/home/ubuntu/training/datasets/indic-asr-benchmark-6k")
    parser.add_argument("--output-dir", default="/home/ubuntu/training/benchmark_results")
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--max-new-tokens", type=int, default=512)
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument("--backend", choices=["transformers", "vllm"], default="vllm",
                        help="Inference backend (default: vllm)")
    parser.add_argument("--gpu-memory-utilization", type=float, default=0.85,
                        help="vLLM GPU memory utilization (default: 0.85)")
    parser.add_argument("--no-force-language", action="store_true",
                        help="Don't force language, let model auto-detect")
    parser.add_argument("--languages", nargs="+", default=None,
                        help="Filter to specific language codes (e.g. hi en ta)")
    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 qwen_asr import Qwen3ASRModel

    print(f"Loading checkpoint: {args.checkpoint} (backend={args.backend})")
    t0 = time.time()
    if args.backend == "vllm":
        model = Qwen3ASRModel.LLM(
            model=args.checkpoint,
            gpu_memory_utilization=args.gpu_memory_utilization,
            max_inference_batch_size=args.batch_size if args.batch_size > 0 else -1,
            max_new_tokens=args.max_new_tokens,
        )
    else:
        model = Qwen3ASRModel.from_pretrained(
            args.checkpoint,
            dtype=torch.bfloat16,
            device_map=args.device,
            max_inference_batch_size=args.batch_size,
            max_new_tokens=args.max_new_tokens,
        )
    load_time = time.time() - t0
    print(f"Model loaded in {load_time:.1f}s")

    # Load dataset
    hf_token = args.hf_token or 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"
    )

    # Filter languages if specified
    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}, "
          f"force_language={not args.no_force_language})...")
    t0 = time.time()
    results = run_inference(model, dataset, args.batch_size,
                            force_language=not args.no_force_language)
    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 metrics
    metrics = compute_metrics(results)
    metrics["__meta__"] = {
        "checkpoint": args.checkpoint,
        "checkpoint_name": ckpt_name,
        "backend": args.backend,
        "dataset": args.hf_dataset,
        "batch_size": args.batch_size,
        "force_language": not args.no_force_language,
        "inference_time_sec": round(inference_time, 2),
        "total_audio_sec": round(total_audio_sec, 2),
        "rtf": round(rtf, 4),
    }

    metrics_file = output_dir / "metrics.json"
    with open(metrics_file, "w") as f:
        json.dump(metrics, f, indent=2, ensure_ascii=False)
    print(f"Metrics saved: {metrics_file}")

    # Print table
    print_metrics_table(metrics, ckpt_name)

    # Save sample analysis
    save_sample_analysis(results, output_dir)

if __name__ == "__main__":
    main()

+# vLLM v1 (0.16+) spawns a subprocess for the engine — CUDA must not be

+# initialized in the parent before that fork.  Setting the env var up-front

+# tells vLLM to use 'spawn' instead of 'fork', which avoids the

+# "Cannot re-initialize CUDA in forked subprocess" crash.


Now defer `import torch` to only when the transformers backend is used (vLLM handles its own torch internally):


    # Load model
    from qwen_asr import Qwen3ASRModel

    print(f"Loading checkpoint: {args.checkpoint} (backend={args.backend})")
    t0 = time.time()
    if args.backend == "vllm":
        model = Qwen3ASRModel.LLM(
            model=args.checkpoint,
            gpu_memory_utilization=args.gpu_memory_utilization,
            max_inference_batch_size=args.batch_size if args.batch_size > 0 else -1,
            max_new_tokens=args.max_new_tokens,
        )
    else:
        model = Qwen3ASRModel.from_pretrained(
            args.checkpoint,
            dtype=torch.bfloat16,
            device_map=args.device,
            max_inference_batch_size=args.batch_size,
            max_new_tokens=args.max_new_tokens,
        )

    # Load model
    from qwen_asr import Qwen3ASRModel

    print(f"Loading checkpoint: {args.checkpoint} (backend={args.backend})")
    t0 = time.time()
    if args.backend == "vllm":
        model = Qwen3ASRModel.LLM(
            model=args.checkpoint,
            gpu_memory_utilization=args.gpu_memory_utilization,
            max_inference_batch_size=args.batch_size if args.batch_size > 0 else -1,
            max_new_tokens=args.max_new_tokens,
        )
    else:
        import torch
        model = Qwen3ASRModel.from_pretrained(
            args.checkpoint,
            dtype=torch.bfloat16,
            device_map=args.device,
            max_inference_batch_size=args.batch_size,
            max_new_tokens=args.max_new_tokens,
        )