  [6000/6000] (100%) processed
Inference done in 181.6s for 40354s audio (RTF=0.0045)
Predictions saved: /home/ubuntu/training/benchmark_outputs/cohere-transcribe/ckpt-10000/predictions.json

Computing metrics...

=====================================================================================
  BENCHMARK RESULTS: ckpt-10000 (cohere-transcribe)
=====================================================================================
Language           N  WER-raw    WER-n   WER-nc      SNW      MER    CER-n
-------------- ----- -------- -------- -------- -------- -------- --------
assamese         500   147.44   147.44   147.44   100.00   124.07   115.39
bengali          500   146.57   146.77   146.77   100.00   114.08   108.05
english          500    61.52    54.20    54.00    45.84    36.36    35.16
gujarati         500   166.14   166.14   166.14   100.00   135.56   128.92
hindi            500   126.93   126.93   126.93   100.00   122.82   111.08
kannada          500   209.06   209.06   209.06   100.00   124.36   122.65
malayalam        500   212.76   212.76   212.76   100.00   115.50   112.97
marathi          500   158.00   158.00   158.00   100.00   119.23   113.95
odia             500   178.02   178.00   178.00   100.00   129.36   124.92
punjabi          500   134.52   134.49   134.49   100.00   124.63   113.39
tamil            500   188.38   188.38   188.38   100.00   117.61   113.31
telugu           500   194.49   194.49   194.49   100.00   125.24   122.68
-------------- ----- -------- -------- -------- -------- -------- --------
OVERALL         6000   159.79   159.25   159.23    95.75   116.91   111.67
MACRO-AVG         12   160.32   159.72   159.70    95.49   115.74   110.21
=====================================================================================

Outputs saved to: /home/ubuntu/training/benchmark_outputs/cohere-transcribe/ckpt-10000
  metrics.json, sample_analysis.json, error_analysis.json


Benchmark complete. Let me verify the output files are correctly placed for the dashboard.


total 12M
-rw-rw-r-- 1 ubuntu ubuntu  59K Apr  5 20:12 error_analysis.json
-rw-rw-r-- 1 ubuntu ubuntu 5.3K Apr  5 20:12 metrics.json
-rw-rw-r-- 1 ubuntu ubuntu 3.0M Apr  5 20:12 predictions.json
-rw-rw-r-- 1 ubuntu ubuntu 8.8M Apr  5 20:12 sample_analysis.json


Bash

python3 -c "
import json
m = json.load(open('/home/ubuntu/training/benchmark_outputs/cohere-transcribe/ckpt-10000/metrics.json'))
# Verify schema compliance
assert '__overall__' in m
assert '__macro_avg__' in m
assert '__meta__' in m
meta = m['__meta__']
assert meta['model_id'] == 'cohere-transcribe'
assert meta['normalization_version'] == 'v1'
print('Schema validation passed')
print(f'Model: {meta[\"model_id\"]}')
print(f'Checkpoint: {meta[\"checkpoint_name\"]}')
print(f'Inference time: {meta[\"inference_time_sec\"]}s')
print(f'RTF: {meta[\"rtf\"]}')
print(f'Languages: {m[\"__macro_avg__\"][\"n_languages\"]}')
"


Schema validation passed
Model: cohere-transcribe
Checkpoint: ckpt-10000
Inference time: 181.56s
RTF: 0.0045
Languages: 12


Updated task #4 status


Benchmark complete. All 3 schema-v1 files are in `benchmark_outputs/cohere-transcribe/ckpt-10000/` and ready for the dashboard.

**Results summary (ckpt-10000):**

| Metric | Overall | English | Hindi |
|--------|---------|---------|-------|
| WER-norm | 159.25% | 54.20% | 126.93% |
| MER | 116.91% | 36.36% | 122.82% |
| CER-norm | 111.67% | 35.16% | 111.08% |
| RTF | 0.0045 | — | — |

At 10k steps, the model still defaults to English transcription for all Indic audio — WER >100% across all 11 Indic languages (lots of insertions from English hallucinations). English itself is at 54% WER-norm, showing the base model's English capability is intact but degraded by the tokenizer extension. This is expected early in fine-tuning; the Indic language conditioning should improve dramatically with more training steps.

**Key note for future checkpoints:** The benchmark script (`benchmark_cohere_transcribe.py`) handles the safetensors key remapping (`encoder_encoder_decoder_proj` → `encoder_decoder_proj`, `transf_decoder._transf_decoder._decoder` → `transf_decoder._decoder`). Just point it at the next checkpoint and it'll work.


are the detected languages correct ? are they all detected correctly by the model ?


Good catch — my code is **not** actually detecting languages from the model output. I'm just passing through the ground truth `lang_code` as `detected_language`:

```python
results.append({**meta, "hypothesis": hyp, "detected_language": meta["lang_code"]})
```

So every sample says "detected: hi" or "detected: ta" etc., but the model is clearly **outputting English** for all Indic audio. The `detected_language` field is misleading.

The Cohere Transcribe model uses a forced language prompt (no auto-detect mode), so there's no built-in language detection. But I can add a simple script-based detection from the hypothesis text to flag what language the model actually produced.


python3 -c "
import json, unicodedata, re
from collections import Counter

# Simple script detection
def detect_script(text):
    scripts = Counter()
    for c in text:
        if c.isalpha():
            cp = ord(c)
            if 0x0900 <= cp <= 0x097F: scripts['Devanagari'] += 1
            elif 0x0980 <= cp <= 0x09FF: scripts['Bengali'] += 1
            elif 0x0A00 <= cp <= 0x0A7F: scripts['Gurmukhi'] += 1
            elif 0x0A80 <= cp <= 0x0AFF: scripts['Gujarati'] += 1
            elif 0x0B00 <= cp <= 0x0B7F: scripts['Odia'] += 1
            elif 0x0B80 <= cp <= 0x0BFF: scripts['Tamil'] += 1
            elif 0x0C00 <= cp <= 0x0C7F: scripts['Telugu'] += 1
            elif 0x0C80 <= cp <= 0x0CFF: scripts['Kannada'] += 1
            elif 0x0D00 <= cp <= 0x0D7F: scripts['Malayalam'] += 1
            elif 0x0041 <= cp <= 0x007A: scripts['Latin'] += 1
    return scripts.most_common(1)[0][0] if scripts else 'Unknown'

preds = json.load(open('/home/ubuntu/training/benchmark_outputs/cohere-transcribe/ckpt-10000/predictions.json'))

# Count actual output script per expected language
by_lang = {}
for p in preds:
    lang = p['language']
    if lang not in by_lang:
        by_lang[lang] = Counter()
    hyp = p['hypothesis']
    if hyp.strip():
        script = detect_script(hyp)
        by_lang[lang][script] += 1
    else:
        by_lang[lang]['Empty'] += 1

print(f'{\"Language\":<12} {\"Expected Script\":<14} {\"Actual Output Scripts\":<50}')
print('-' * 76)
for lang in sorted(by_lang):
    expected = {'hindi': 'Devanagari', 'bengali': 'Bengali', 'tamil': 'Tamil',
                'telugu': 'Telugu', 'kannada': 'Kannada', 'malayalam': 'Malayalam',
                'gujarati': 'Gujarati', 'punjabi': 'Gurmukhi', 'marathi': 'Devanagari',
                'odia': 'Odia', 'assamese': 'Bengali', 'english': 'Latin'}[lang]
    counts = by_lang[lang].most_common()
    count_str = ', '.join(f'{s}:{n}' for s, n in counts)
    print(f'{lang:<12} {expected:<14} {count_str}')
"


Language     Expected Script Actual Output Scripts
----------------------------------------------------------------------------
assamese     Bengali        Latin:491, Empty:9
bengali      Bengali        Latin:500
english      Latin          Latin:491, Empty:9
gujarati     Gujarati       Latin:500
hindi        Devanagari     Latin:500
kannada      Kannada        Latin:500
malayalam    Malayalam      Latin:500
marathi      Devanagari     Latin:500
odia         Odia           Latin:500
punjabi      Gurmukhi       Latin:500
tamil        Tamil          Latin:500
telugu       Telugu         Latin:500


Every single Indic language output is in Latin script (English). The model at ckpt-10000 is completely ignoring the `<|hi|>`, `<|ta|>` etc. language tokens in the decoder prompt and defaulting to English transcription for everything. This is expected — at only 10k training steps, the Indic decoder heads / language conditioning haven't learned yet.

The `detected_language` field in the output files is currently just the ground truth label, which is incorrect. Let me fix it to reflect the actual output script.


300	    return analysis
301
302
303	# ── Inference ────────────────────────────────────────────────────────────
304
305	def run_inference(model, processor, tokenizer, dataset, batch_size=1, max_new_tokens=256, device="cuda:0"):
306	    """Run batched inference on HF dataset."""
307	    results = []
308	    total = len(dataset)
309
310	    # Pre-cache prompt IDs per language code
311	    prompt_cache = {}
312	    for lc in LANG_CODE_TO_NAME:
313	        prompt_str = model.build_prompt(language=lc, punctuation=True)
314	        prompt_cache[lc] = tokenizer.encode(prompt_str, add_special_tokens=False)

    return analysis

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

def run_inference(model, processor, tokenizer, dataset, batch_size=1, max_new_tokens=256, device="cuda:0"):
    """Run batched inference on HF dataset."""
    results = []
    total = len(dataset)

    # Pre-cache prompt IDs per language code
    prompt_cache = {}
    for lc in LANG_CODE_TO_NAME:
        prompt_str = model.build_prompt(language=lc, punctuation=True)
        prompt_cache[lc] = tokenizer.encode(prompt_str, add_special_tokens=False)


Edit

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

def run_inference(model, processor, tokenizer, dataset, batch_size=1, max_new_tokens=256, device="cuda:0"):
    """Run batched inference on HF dataset."""

# ── Script detection ─────────────────────────────────────────────────────

SCRIPT_RANGES = [
    (0x0900, 0x097F, "Devanagari"), (0x0980, 0x09FF, "Bengali"),
    (0x0A00, 0x0A7F, "Gurmukhi"), (0x0A80, 0x0AFF, "Gujarati"),
    (0x0B00, 0x0B7F, "Odia"), (0x0B80, 0x0BFF, "Tamil"),
    (0x0C00, 0x0C7F, "Telugu"), (0x0C80, 0x0CFF, "Kannada"),
    (0x0D00, 0x0D7F, "Malayalam"), (0x0041, 0x007A, "Latin"),
]

SCRIPT_TO_LANG = {
    "Devanagari": "Hindi", "Bengali": "Bengali", "Gurmukhi": "Punjabi",
    "Gujarati": "Gujarati", "Odia": "Odia", "Tamil": "Tamil",
    "Telugu": "Telugu", "Kannada": "Kannada", "Malayalam": "Malayalam",
    "Latin": "English",
}

def detect_output_language(text: str) -> str:
    """Detect the dominant script/language of model output text."""
    from collections import Counter
    scripts = Counter()
    for c in text:
        if c.isalpha():
            cp = ord(c)
            for lo, hi, name in SCRIPT_RANGES:
                if lo <= cp <= hi:
                    scripts[name] += 1
                    break
    if not scripts:
        return ""
    dominant = scripts.most_common(1)[0][0]
    return SCRIPT_TO_LANG.get(dominant, dominant)

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

def run_inference(model, processor, tokenizer, dataset, batch_size=1, max_new_tokens=256, device="cuda:0"):
    """Run batched inference on HF dataset."""


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

# ── Script detection ─────────────────────────────────────────────────────

SCRIPT_RANGES = [
    (0x0900, 0x097F, "Devanagari"), (0x0980, 0x09FF, "Bengali"),
    (0x0A00, 0x0A7F, "Gurmukhi"), (0x0A80, 0x0AFF, "Gujarati"),
    (0x0B00, 0x0B7F, "Odia"), (0x0B80, 0x0BFF, "Tamil"),
    (0x0C00, 0x0C7F, "Telugu"), (0x0C80, 0x0CFF, "Kannada"),
    (0x0D00, 0x0D7F, "Malayalam"), (0x0041, 0x007A, "Latin"),
]

SCRIPT_TO_LANG = {
    "Devanagari": "Hindi", "Bengali": "Bengali", "Gurmukhi": "Punjabi",
    "Gujarati": "Gujarati", "Odia": "Odia", "Tamil": "Tamil",
    "Telugu": "Telugu", "Kannada": "Kannada", "Malayalam": "Malayalam",
    "Latin": "English",
}

def detect_output_language(text: str) -> str:
    """Detect the dominant script/language of model output text."""
    from collections import Counter
    scripts = Counter()
    for c in text:
        if c.isalpha():
            cp = ord(c)
            for lo, hi, name in SCRIPT_RANGES:
                if lo <= cp <= hi:
                    scripts[name] += 1
                    break
    if not scripts:
        return ""
    dominant = scripts.most_common(1)[0][0]
    return SCRIPT_TO_LANG.get(dominant, dominant)

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

def run_inference(model, processor, tokenizer, dataset, batch_size=1, max_new_tokens=256, device="cuda:0"):
    """Run batched inference on HF dataset."""

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

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

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

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

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

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

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

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

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

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

LATIN_SCRIPT_LANGS = {"english"}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    return (len(touched), total_words)

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

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

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

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

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

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

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

        wer_scores = []

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

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

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

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

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

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

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

    return analysis

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

def run_inference(model, processor, tokenizer, dataset, batch_size=1, max_new_tokens=256, device="cuda:0"):
    """Run batched inference on HF dataset."""
    results = []
    total = len(dataset)

    # Pre-cache prompt IDs per language code
    prompt_cache = {}
    for lc in LANG_CODE_TO_NAME:
        prompt_str = model.build_prompt(language=lc, punctuation=True)
        prompt_cache[lc] = tokenizer.encode(prompt_str, add_special_tokens=False)

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

        # Collect audio arrays
        wavs = []
        batch_meta = []
        for sample in batch:
            audio = sample["audio"]
            wav = np.array(audio["array"], dtype=np.float32)
            sr = audio["sampling_rate"]
            if sr != 16000:
                wav_t = torch.from_numpy(wav).unsqueeze(0)
                wav_t = torchaudio.functional.resample(wav_t, sr, 16000)
                wav = wav_t.squeeze(0).numpy()
            wavs.append(wav)
            batch_meta.append({
                "id": sample["id"],
                "language": sample["language"],
                "lang_code": sample["lang_code"],
                "source": sample["source"],
                "duration": sample["duration"],
                "reference": sample["reference"],
            })

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

        # Build batched decoder prompts (pad to same length — all prompts are same length)
        prompt_ids_list = [prompt_cache[m["lang_code"]] for m in batch_meta]
        max_prompt_len = max(len(p) for p in prompt_ids_list)
        pad_id = tokenizer.convert_tokens_to_ids("<pad>")
        padded_prompts = [p + [pad_id] * (max_prompt_len - len(p)) for p in prompt_ids_list]
        decoder_input_ids = torch.tensor(padded_prompts, device=device)

        # Batched generate
        with torch.no_grad():
            try:
                outputs = model.generate(
                    input_features=input_features,
                    length=length,
                    decoder_input_ids=decoder_input_ids,
                    max_new_tokens=max_new_tokens,
                )
                for i, meta in enumerate(batch_meta):
                    hyp = decode_tokens(tokenizer, outputs[i]).strip()
                    results.append({**meta, "hypothesis": hyp, "detected_language": meta["lang_code"]})
            except Exception as e:
                print(f"  Batch {start}-{end} failed ({e}), falling back to one-by-one...", flush=True)
                for i, meta in enumerate(batch_meta):
                    try:
                        feat_i = input_features[i:i+1]
                        len_i = length[i:i+1]
                        did_i = torch.tensor([prompt_ids_list[i]], device=device)
                        out_i = model.generate(
                            input_features=feat_i, length=len_i,
                            decoder_input_ids=did_i, max_new_tokens=max_new_tokens,
                        )
                        hyp = decode_tokens(tokenizer, out_i[0]).strip()
                    except Exception as e2:
                        print(f"  Sample {meta['id']} failed: {e2}")
                        hyp = ""
                    results.append({**meta, "hypothesis": hyp, "detected_language": meta["lang_code"]})

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

    return results

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

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

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

    # Load model
    from transformers import AutoModelForSpeechSeq2Seq, AutoFeatureExtractor

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

    # Load feature extractor from checkpoint (has custom processing code)
    feature_extractor = AutoFeatureExtractor.from_pretrained(
        args.checkpoint, trust_remote_code=True,
    )
    # FilterbankFeatures is lazily created — set device in both _device and
    # _fb_config so the torch.Generator inside FilterbankFeatures is on GPU
    feature_extractor._device = args.device
    feature_extractor._fb_config["device"] = args.device

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

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

    # Load model — checkpoint has different key prefixes from training wrapper,
    # so we load the architecture first, then manually load+remap weights.
    from safetensors.torch import load_file as load_safetensors

    model = AutoModelForSpeechSeq2Seq.from_pretrained(
        args.checkpoint,
        trust_remote_code=True,
        torch_dtype=torch.bfloat16,
        ignore_mismatched_sizes=True,
    )

    # Remap checkpoint keys to match model's expected names
    ckpt_weights = load_safetensors(str(Path(args.checkpoint) / "model.safetensors"))
    KEY_REMAP = [
        ("encoder_encoder_decoder_proj.", "encoder_decoder_proj."),
        ("transf_decoder._transf_decoder._decoder.", "transf_decoder._decoder."),
    ]
    remapped = {}
    for k, v in ckpt_weights.items():
        new_k = k
        for old_prefix, new_prefix in KEY_REMAP:
            if new_k.startswith(old_prefix):
                new_k = new_prefix + new_k[len(old_prefix):]
        remapped[new_k] = v

    missing, unexpected = model.load_state_dict(remapped, strict=False)
    if missing:
        print(f"WARNING: {len(missing)} missing keys (first 5): {missing[:5]}")
    if unexpected:
        print(f"WARNING: {len(unexpected)} unexpected keys (first 5): {unexpected[:5]}")

    model = model.to(args.device)
    model.eval()

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    error_analysis = compute_error_analysis(samples_for_errors)

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

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

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

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

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

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

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

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

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

if __name__ == "__main__":
    main()

+# ── Script detection ─────────────────────────────────────────────────────

+    "Devanagari": "Hindi", "Bengali": "Bengali", "Gurmukhi": "Punjabi",

+    "Gujarati": "Gujarati", "Odia": "Odia", "Tamil": "Tamil",

+    "Telugu": "Telugu", "Kannada": "Kannada", "Malayalam": "Malayalam",

+def detect_output_language(text: str) -> str:

+    """Detect the dominant script/language of model output text."""

+            for lo, hi, name in SCRIPT_RANGES:

+    dominant = scripts.most_common(1)[0][0]

+    return SCRIPT_TO_LANG.get(dominant, dominant)


Now update the results to use the actual detected language instead of the ground truth.


400:                    results.append({**meta, "hypothesis": hyp, "detected_language": meta["lang_code"]})
416:                    results.append({**meta, "hypothesis": hyp, "detected_language": meta["lang_code"]})


                    results.append({**meta, "hypothesis": hyp, "detected_language": meta["lang_code"]})
            except Exception as e:
                print(f"  Batch {start}-{end} failed ({e}), falling back to one-by-one...", flush=True)
                for i, meta in enumerate(batch_meta):
                    try:
                        feat_i = input_features[i:i+1]
                        len_i = length[i:i+1]
                        did_i = torch.tensor([prompt_ids_list[i]], device=device)
                        out_i = model.generate(
                            input_features=feat_i, length=len_i,
                            decoder_input_ids=did_i, max_new_tokens=max_new_tokens,
                        )
                        hyp = decode_tokens(tokenizer, out_i[0]).strip()
                    except Exception as e2:
                        print(f"  Sample {meta['id']} failed: {e2}")
                        hyp = ""
                    results.append({**meta, "hypothesis": hyp, "detected_language": meta["lang_code"]})

                    detected = detect_output_language(hyp)
                    results.append({**meta, "hypothesis": hyp, "detected_language": detected})
            except Exception as e:
                print(f"  Batch {start}-{end} failed ({e}), falling back to one-by-one...", flush=True)
                for i, meta in enumerate(batch_meta):
                    try:
                        feat_i = input_features[i:i+1]
                        len_i = length[i:i+1]
                        did_i = torch.tensor([prompt_ids_list[i]], device=device)
                        out_i = model.generate(
                            input_features=feat_i, length=len_i,
                            decoder_input_ids=did_i, max_new_tokens=max_new_tokens,
                        )
                        hyp = decode_tokens(tokenizer, out_i[0]).strip()
                    except Exception as e2:
                        print(f"  Sample {meta['id']} failed: {e2}")
                        hyp = ""
                    detected = detect_output_language(hyp)
                    results.append({**meta, "hypothesis": hyp, "detected_language": detected})


                    detected = detect_output_language(hyp)
                    results.append({**meta, "hypothesis": hyp, "detected_language": detected})
            except Exception as e:
                print(f"  Batch {start}-{end} failed ({e}), falling back to one-by-one...", flush=True)
                for i, meta in enumerate(batch_meta):
                    try:
                        feat_i = input_features[i:i+1]
                        len_i = length[i:i+1]
                        did_i = torch.tensor([prompt_ids_list[i]], device=device)
                        out_i = model.generate(
                            input_features=feat_i, length=len_i,
                            decoder_input_ids=did_i, max_new_tokens=max_new_tokens,
                        )
                        hyp = decode_tokens(tokenizer, out_i[0]).strip()
                    except Exception as e2:
                        print(f"  Sample {meta['id']} failed: {e2}")
                        hyp = ""
                    detected = detect_output_language(hyp)
                    results.append({**meta, "hypothesis": hyp, "detected_language": detected})

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

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

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

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

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

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

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

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

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

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

LATIN_SCRIPT_LANGS = {"english"}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    return (len(touched), total_words)

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

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

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

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

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

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

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

        wer_scores = []

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

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

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

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

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

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

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

    return analysis

# ── Script detection ─────────────────────────────────────────────────────

SCRIPT_RANGES = [
    (0x0900, 0x097F, "Devanagari"), (0x0980, 0x09FF, "Bengali"),
    (0x0A00, 0x0A7F, "Gurmukhi"), (0x0A80, 0x0AFF, "Gujarati"),
    (0x0B00, 0x0B7F, "Odia"), (0x0B80, 0x0BFF, "Tamil"),
    (0x0C00, 0x0C7F, "Telugu"), (0x0C80, 0x0CFF, "Kannada"),
    (0x0D00, 0x0D7F, "Malayalam"), (0x0041, 0x007A, "Latin"),
]

SCRIPT_TO_LANG = {
    "Devanagari": "Hindi", "Bengali": "Bengali", "Gurmukhi": "Punjabi",
    "Gujarati": "Gujarati", "Odia": "Odia", "Tamil": "Tamil",
    "Telugu": "Telugu", "Kannada": "Kannada", "Malayalam": "Malayalam",
    "Latin": "English",
}

def detect_output_language(text: str) -> str:
    """Detect the dominant script/language of model output text."""
    from collections import Counter
    scripts = Counter()
    for c in text:
        if c.isalpha():
            cp = ord(c)
            for lo, hi, name in SCRIPT_RANGES:
                if lo <= cp <= hi:
                    scripts[name] += 1
                    break
    if not scripts:
        return ""
    dominant = scripts.most_common(1)[0][0]
    return SCRIPT_TO_LANG.get(dominant, dominant)

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

def run_inference(model, processor, tokenizer, dataset, batch_size=1, max_new_tokens=256, device="cuda:0"):
    """Run batched inference on HF dataset."""
    results = []
    total = len(dataset)

    # Pre-cache prompt IDs per language code
    prompt_cache = {}
    for lc in LANG_CODE_TO_NAME:
        prompt_str = model.build_prompt(language=lc, punctuation=True)
        prompt_cache[lc] = tokenizer.encode(prompt_str, add_special_tokens=False)

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

        # Collect audio arrays
        wavs = []
        batch_meta = []
        for sample in batch:
            audio = sample["audio"]
            wav = np.array(audio["array"], dtype=np.float32)
            sr = audio["sampling_rate"]
            if sr != 16000:
                wav_t = torch.from_numpy(wav).unsqueeze(0)
                wav_t = torchaudio.functional.resample(wav_t, sr, 16000)
                wav = wav_t.squeeze(0).numpy()
            wavs.append(wav)
            batch_meta.append({
                "id": sample["id"],
                "language": sample["language"],
                "lang_code": sample["lang_code"],
                "source": sample["source"],
                "duration": sample["duration"],
                "reference": sample["reference"],
            })

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

        # Build batched decoder prompts (pad to same length — all prompts are same length)
        prompt_ids_list = [prompt_cache[m["lang_code"]] for m in batch_meta]
        max_prompt_len = max(len(p) for p in prompt_ids_list)
        pad_id = tokenizer.convert_tokens_to_ids("<pad>")
        padded_prompts = [p + [pad_id] * (max_prompt_len - len(p)) for p in prompt_ids_list]
        decoder_input_ids = torch.tensor(padded_prompts, device=device)

        # Batched generate
        with torch.no_grad():
            try:
                outputs = model.generate(
                    input_features=input_features,
                    length=length,
                    decoder_input_ids=decoder_input_ids,
                    max_new_tokens=max_new_tokens,
                )
                for i, meta in enumerate(batch_meta):
                    hyp = decode_tokens(tokenizer, outputs[i]).strip()
                    results.append({**meta, "hypothesis": hyp, "detected_language": meta["lang_code"]})
            except Exception as e:
                print(f"  Batch {start}-{end} failed ({e}), falling back to one-by-one...", flush=True)
                for i, meta in enumerate(batch_meta):
                    try:
                        feat_i = input_features[i:i+1]
                        len_i = length[i:i+1]
                        did_i = torch.tensor([prompt_ids_list[i]], device=device)
                        out_i = model.generate(
                            input_features=feat_i, length=len_i,
                            decoder_input_ids=did_i, max_new_tokens=max_new_tokens,
                        )
                        hyp = decode_tokens(tokenizer, out_i[0]).strip()
                    except Exception as e2:
                        print(f"  Sample {meta['id']} failed: {e2}")
                        hyp = ""
                    results.append({**meta, "hypothesis": hyp, "detected_language": meta["lang_code"]})

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

    return results

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

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

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

    # Load model
    from transformers import AutoModelForSpeechSeq2Seq, AutoFeatureExtractor

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

    # Load feature extractor from checkpoint (has custom processing code)
    feature_extractor = AutoFeatureExtractor.from_pretrained(
        args.checkpoint, trust_remote_code=True,
    )
    # FilterbankFeatures is lazily created — set device in both _device and
    # _fb_config so the torch.Generator inside FilterbankFeatures is on GPU
    feature_extractor._device = args.device
    feature_extractor._fb_config["device"] = args.device

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

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

    # Load model — checkpoint has different key prefixes from training wrapper,
    # so we load the architecture first, then manually load+remap weights.
    from safetensors.torch import load_file as load_safetensors

    model = AutoModelForSpeechSeq2Seq.from_pretrained(
        args.checkpoint,
        trust_remote_code=True,
        torch_dtype=torch.bfloat16,
        ignore_mismatched_sizes=True,
    )

    # Remap checkpoint keys to match model's expected names
    ckpt_weights = load_safetensors(str(Path(args.checkpoint) / "model.safetensors"))
    KEY_REMAP = [
        ("encoder_encoder_decoder_proj.", "encoder_decoder_proj."),
        ("transf_decoder._transf_decoder._decoder.", "transf_decoder._decoder."),
    ]
    remapped = {}
    for k, v in ckpt_weights.items():
        new_k = k
        for old_prefix, new_prefix in KEY_REMAP:
            if new_k.startswith(old_prefix):
                new_k = new_prefix + new_k[len(old_prefix):]
        remapped[new_k] = v

    missing, unexpected = model.load_state_dict(remapped, strict=False)
    if missing:
        print(f"WARNING: {len(missing)} missing keys (first 5): {missing[:5]}")
    if unexpected:
        print(f"WARNING: {len(unexpected)} unexpected keys (first 5): {unexpected[:5]}")

    model = model.to(args.device)
    model.eval()

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    error_analysis = compute_error_analysis(samples_for_errors)

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

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

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

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

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

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

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

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

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

if __name__ == "__main__":
    main()

-                    results.append({**meta, "hypothesis": hyp, "detected_language": meta["lang_code"]})

+                    detected = detect_output_language(hyp)

+                    results.append({**meta, "hypothesis": hyp, "detected_language": detected})


Now re-run the full benchmark with corrected `detected_language`.


Bash

python3 /home/ubuntu/training/benchmark_cohere_transcribe.py \
    --checkpoint /home/ubuntu/training/checkpoints/cohere-transcribe-ckpt-10000 \
    --checkpoint-name ckpt-10000 \
    --batch-size 64 \
    --device cuda:0 2>&1 | tail -25

Full benchmark with detected_language fix


=====================================================================================
  BENCHMARK RESULTS: ckpt-10000 (cohere-transcribe)
=====================================================================================
Language           N  WER-raw    WER-n   WER-nc      SNW      MER    CER-n
-------------- ----- -------- -------- -------- -------- -------- --------
assamese         500   147.44   147.44   147.44   100.00   124.07   115.39
bengali          500   146.57   146.77   146.77   100.00   114.08   108.05
english          500    61.52    54.20    54.00    45.84    36.36    35.16
gujarati         500   166.14   166.14   166.14   100.00   135.56   128.92
hindi            500   126.93   126.93   126.93   100.00   122.82   111.08
kannada          500   209.06   209.06   209.06   100.00   124.36   122.65
malayalam        500   212.76   212.76   212.76   100.00   115.50   112.97
marathi          500   158.00   158.00   158.00   100.00   119.23   113.95
odia             500   178.02   178.00   178.00   100.00   129.36   124.92
punjabi          500   134.52   134.49   134.49   100.00   124.63   113.39
tamil            500   188.38   188.38   188.38   100.00   117.61   113.31
telugu           500   194.49   194.49   194.49   100.00   125.24   122.68
-------------- ----- -------- -------- -------- -------- -------- --------
OVERALL         6000   159.79   159.25   159.23    95.75   116.91   111.67
MACRO-AVG         12   160.32   159.72   159.70    95.49   115.74   110.21
=====================================================================================

Outputs saved to: /home/ubuntu/training/benchmark_outputs/cohere-transcribe/ckpt-10000
  metrics.json, sample_analysis.json, error_analysis.json


Let me verify the `detected_language` is now correct in the output.


python3 -c "
import json
from collections import Counter

preds = json.load(open('/home/ubuntu/training/benchmark_outputs/cohere-transcribe/ckpt-10000/predictions.json'))
by_lang = {}
for p in preds:
    lang = p['language']
    if lang not in by_lang:
        by_lang[lang] = Counter()
    by_lang[lang][p['detected_language']] += 1

print(f'{\"Expected\":<12} {\"Detected output language\":<50}')
print('-' * 62)
for lang in sorted(by_lang):
    counts = by_lang[lang].most_common()
    count_str = ', '.join(f'{s}:{n}' for s, n in counts)
    print(f'{lang:<12} {count_str}')
"


Expected     Detected output language
--------------------------------------------------------------
assamese     English:491, :9
bengali      English:500
english      English:491, :9
gujarati     English:500
hindi        English:500
kannada      English:500
malayalam    English:500
marathi      English:500
odia         English:500
punjabi      English:500
tamil        English:500
telugu       English:500


Now the `detected_language` field correctly reflects reality: the model outputs English for everything at ckpt-10000. The `lang_confusion` flag will also fire in the sample_analysis for all Indic samples.


python3 -c "
import json
sa = json.load(open('/home/ubuntu/training/benchmark_outputs/cohere-transcribe/ckpt-10000/sample_analysis.json'))
from collections import Counter
flags = Counter()
for s in sa:
    for f in s['flags']:
        flags[f] += 1
for f, c in flags.most_common():
    print(f'{f:<25} {c:>5}')
"


high_wer                   5669
spacing_error               396
numeric_mismatch             62
exact_match_norm             41
punctuation_only_diff        22
exact_match                  19
empty_hypothesis             18


The `lang_confusion` flag isn't firing — I need to check the flag detection. The issue is my `detect_flags` function doesn't compare `detected_language` with expected language. Let me check.


192:def detect_flags(ref: str, hyp: str, ref_n: str, hyp_n: str,
401:                    results.append({**meta, "hypothesis": hyp, "detected_language": detected})
418:                    results.append({**meta, "hypothesis": hyp, "detected_language": detected})
582:            "detected_language": p.get("detected_language", ""),
693:            "detected_language": s["detected_language"],


192	def detect_flags(ref: str, hyp: str, ref_n: str, hyp_n: str,
193	                 wer_norm_val: float, mer_val: float) -> list:
194	    flags = []
195	    if ref == hyp:
196	        flags.append("exact_match")
197	    if ref_n == hyp_n:
198	        flags.append("exact_match_norm")
199	    if not hyp.strip():
200	        flags.append("empty_hypothesis")
201	    if ref != hyp and ref_n == hyp_n:
202	        flags.append("punctuation_only_diff")
203	    ref_digits = re.findall(r"\d+", ref_n)
204	    hyp_digits = re.findall(r"\d+", hyp_n)
205	    if ref_digits != hyp_digits and (ref_digits or hyp_digits):
206	        flags.append("numeric_mismatch")
207	    if wer_norm_val > 80:
208	        flags.append("high_wer")
209	    if wer_norm_val > mer_val and mer_val < 100:
210	        flags.append("spacing_error")
211	    return flags

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


Right — `detect_flags` doesn't take `detected_language` or `expected_language`. Let me add `lang_confusion` and `script_mismatch` flags.


Edit

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

def detect_flags(ref: str, hyp: str, ref_n: str, hyp_n: str,
                 wer_norm_val: float, mer_val: float,
                 detected_language: str = "", expected_language: str = "") -> list:
    flags = []
    if ref == hyp:
        flags.append("exact_match")
    if ref_n == hyp_n:
        flags.append("exact_match_norm")
    if not hyp.strip():
        flags.append("empty_hypothesis")
    if ref != hyp and ref_n == hyp_n:
        flags.append("punctuation_only_diff")
    ref_digits = re.findall(r"\d+", ref_n)
    hyp_digits = re.findall(r"\d+", hyp_n)
    if ref_digits != hyp_digits and (ref_digits or hyp_digits):
        flags.append("numeric_mismatch")
    if wer_norm_val > 80:
        flags.append("high_wer")
    if wer_norm_val > mer_val and mer_val < 100:
        flags.append("spacing_error")
    # Language/script confusion
    if detected_language and expected_language:
        if detected_language.lower() != expected_language.lower():
            flags.append("lang_confusion")
            # Check if it's a script mismatch (e.g. Latin output for Devanagari expected)
            ref_script = detect_output_language(ref)
            hyp_script = detect_output_language(hyp)
            if ref_script and hyp_script and ref_script != hyp_script:
                flags.append("script_mismatch")
    return flags