410	    model = AutoModelForSpeechSeq2Seq.from_pretrained(
411	        args.checkpoint,
412	        trust_remote_code=True,
413	        torch_dtype=torch.bfloat16,
414	    ).to(args.device)
415	    model.eval()
416
417	    load_time = time.time() - t0
418	    print(f"Model loaded in {load_time:.1f}s")
419
420	    # Load dataset
421	    hf_token = os.environ.get("HF_TOKEN")
422	    print(f"Loading dataset: {args.hf_dataset}")
423	    dataset = hf_load_dataset(
424	        args.hf_dataset, token=hf_token, cache_dir=args.cache_dir, split="train"
425	    )
426	    dataset = dataset.cast_column("audio", Audio(sampling_rate=16000, decode=True))
427
428	    if args.languages:
429	        dataset = dataset.filter(lambda x: x["lang_code"] in args.languages)
430	        print(f"Filtered to {len(dataset)} samples for languages: {args.languages}")
431
432	    print(f"Dataset: {len(dataset)} samples")
433
434	    # Run inference

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

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

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

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

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

    # Run inference


395	    feature_extractor = AutoFeatureExtractor.from_pretrained(
396	        args.checkpoint, trust_remote_code=True, device=args.device
397	    )
398
399	    # Load extended tokenizer
400	    tokenizer = load_extended_tokenizer(args.checkpoint)
401
402	    # Create a simple processor-like namespace
403	    class SimpleProcessor:
404	        pass
405	    processor = SimpleProcessor()
406	    processor.feature_extractor = feature_extractor
407	    processor.tokenizer = tokenizer
408
409	    # Load model
410	    model = AutoModelForSpeechSeq2Seq.from_pretrained(
411	        args.checkpoint,
412	        trust_remote_code=True,
413	        torch_dtype=torch.bfloat16,
414	    ).to(args.device)

    feature_extractor = AutoFeatureExtractor.from_pretrained(
        args.checkpoint, trust_remote_code=True, device=args.device
    )

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

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

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


Edit

    # Load feature extractor from checkpoint (has custom processing code)
    feature_extractor = AutoFeatureExtractor.from_pretrained(
        args.checkpoint, trust_remote_code=True, device=args.device
    )

    # Load feature extractor from checkpoint (has custom processing code)
    # Pass device= so the FilterbankFeatures (and its torch.Generator) are on GPU
    feature_extractor = AutoFeatureExtractor.from_pretrained(
        args.checkpoint, trust_remote_code=True, device=args.device
    )
    # Ensure lazy filterbank will init on the right device
    feature_extractor._device = args.device


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

    # Load feature extractor from checkpoint (has custom processing code)
    # Pass device= so the FilterbankFeatures (and its torch.Generator) are on GPU
    feature_extractor = AutoFeatureExtractor.from_pretrained(
        args.checkpoint, trust_remote_code=True, device=args.device
    )
    # Ensure lazy filterbank will init on the right device
    feature_extractor._device = args.device

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

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

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

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

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

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

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

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

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

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

LATIN_SCRIPT_LANGS = {"english"}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    return (len(touched), total_words)

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

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

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

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

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

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

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

        wer_scores = []

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

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

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

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

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

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

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

    return analysis

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

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

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

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

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

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

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

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

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

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

    return results

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

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

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

    # Load model
    from transformers import AutoModelForSpeechSeq2Seq, AutoFeatureExtractor

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

    # Load feature extractor from checkpoint (has custom processing code)
    feature_extractor = AutoFeatureExtractor.from_pretrained(
        args.checkpoint, trust_remote_code=True, device=args.device
    )

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    error_analysis = compute_error_analysis(samples_for_errors)

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

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

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

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

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

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

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

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

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

if __name__ == "__main__":
    main()

+    # Pass device= so the FilterbankFeatures (and its torch.Generator) are on GPU

+    # Ensure lazy filterbank will init on the right device

+    feature_extractor._device = args.device


Bash

python3 benchmark_cohere_transcribe.py \
    --checkpoint /home/ubuntu/training/checkpoints/cohere-transcribe-ckpt-10000 \
    --checkpoint-name ckpt-10000 \
    --languages hi \
    --batch-size 1 2>&1 | tail -30

Smoke test Hindi only with device fix


Found the latest cached dataset configuration 'default' at /home/ubuntu/training/datasets/indic-asr-benchmark-6k/BayAreaBoys___indic-asr-benchmark-6k/default/0.0.0/4d4bc89b7c915a6d80a1efbee3a708d58e688c81 (last modified on Sun Apr  5 19:58:09 2026).
Loading checkpoint: /home/ubuntu/training/checkpoints/cohere-transcribe-ckpt-10000
Model loaded in 6.3s
Loading dataset: BayAreaBoys/indic-asr-benchmark-6k
Filtered to 500 samples for languages: ['hi']
Dataset: 500 samples

Running inference (batch_size=1)...
Traceback (most recent call last):
  File "/home/ubuntu/training/benchmark_cohere_transcribe.py", line 675, in <module>
    main()
  File "/home/ubuntu/training/benchmark_cohere_transcribe.py", line 440, in main
    results = run_inference(
  File "/home/ubuntu/training/benchmark_cohere_transcribe.py", line 327, in run_inference
    features = processor.feature_extractor(
  File "/home/ubuntu/.cache/huggingface/modules/transformers_modules/cohere_hyphen_transcribe_hyphen_ckpt_hyphen_10000/processing_cohere_asr.py", line 436, in __call__
    input_features, length = self.filterbank(audio_tensor, seq_len)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "/home/ubuntu/.cache/huggingface/modules/transformers_modules/cohere_hyphen_transcribe_hyphen_ckpt_hyphen_10000/processing_cohere_asr.py", line 256, in forward
    x = self._apply_dither(x, seq_len_time)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 838, in _fn
    return fn(*args, **kwargs)
  File "/home/ubuntu/.cache/huggingface/modules/transformers_modules/cohere_hyphen_transcribe_hyphen_ckpt_hyphen_10000/processing_cohere_asr.py", line 166, in _apply_dither
    noise = torch.randn(
RuntimeError: Expected a 'cuda' device type for generator but found 'cpu'


The problem is that the `torch.Generator` inside FilterbankFeatures is created with the device from `_fb_config`, and `.to()` doesn't move generators. Let me fix this properly.


Edit

    # Load feature extractor from checkpoint (has custom processing code)
    # Pass device= so the FilterbankFeatures (and its torch.Generator) are on GPU
    feature_extractor = AutoFeatureExtractor.from_pretrained(
        args.checkpoint, trust_remote_code=True, device=args.device
    )
    # Ensure lazy filterbank will init on the right device
    feature_extractor._device = args.device

    # 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 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

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

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

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

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

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

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

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

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

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

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

LATIN_SCRIPT_LANGS = {"english"}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    return (len(touched), total_words)

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

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

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

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

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

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

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

        wer_scores = []

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

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

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

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

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

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

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

    return analysis

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

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

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

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

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

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

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

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

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

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

    return results

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

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

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

    # Load model
    from transformers import AutoModelForSpeechSeq2Seq, AutoFeatureExtractor

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

    # Load feature extractor from checkpoint (has custom processing code)
    # Pass device= so the FilterbankFeatures (and its torch.Generator) are on GPU
    feature_extractor = AutoFeatureExtractor.from_pretrained(
        args.checkpoint, trust_remote_code=True, device=args.device
    )
    # Ensure lazy filterbank will init on the right device
    feature_extractor._device = args.device

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    error_analysis = compute_error_analysis(samples_for_errors)

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

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

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

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

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

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

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

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

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

if __name__ == "__main__":
    main()

-    # Pass device= so the FilterbankFeatures (and its torch.Generator) are on GPU

-        args.checkpoint, trust_remote_code=True, device=args.device

+        args.checkpoint, trust_remote_code=True,

-    # Ensure lazy filterbank will init on the right device

+    # FilterbankFeatures is lazily created — set device in both _device and

+    # _fb_config so the torch.Generator inside FilterbankFeatures is on GPU

+    feature_extractor._fb_config["device"] = args.device


Bash

python3 benchmark_cohere_transcribe.py \
    --checkpoint /home/ubuntu/training/checkpoints/cohere-transcribe-ckpt-10000 \
    --checkpoint-name ckpt-10000 \
    --languages hi \
    --batch-size 1 2>&1 | tail -30

Smoke test Hindi with filterbank device fix


● Yes, let it continue. The agent is right on both points:

  1. The extended_model files are correct — modeling_cohere_asr.py, tokenization_cohere_asr.py, special_tokens_map.json don't
  change during training. Only model.safetensors changes. The copies are fine.
  2. The dither/device fix — it just needs to move the generator to CUDA before inference. Standard pattern is model.to("cuda")
   then the feature extractor's generator follows, or explicitly set device="cuda" when creating the processor. Let the agent
  fix it.

  Let me also push those missing files to R2 so the checkpoint is self-contained:

● Bash(# Push the missing code files to R2 so ckpt-10000 is fully self-contained
      SRC="/workspace/training/tokenizer_extension/extended_model"…)
  ⎿  === R2 model/ ===
          4203 config.json
          1659 configuration_cohere_asr.py
     … +11 lines (ctrl+o to expand)

● R2 checkpoint is now fully self-contained. All files present for both model loading and inference. Let the other agent
  continue with the dither fix.

✻ Baked for 39s

❯ but we changed the tokenizer size right. it didnt change the special tokens ?

● Good question. Let me verify — the special tokens map in the checkpoint should match what our extended tokenizer actually
  uses:

● Bash(python3 -c "
      import json, sys…)
  ⎿  === SPECIAL TOKENS IN EXTENDED TOKENIZER ===
     pad: <pad> (id=2)
     eos: <|endoftext|> (id=3)
     … +39 lines (ctrl+o to expand)

● All special tokens are preserved with their original IDs. The tokenizer extension only added new tokens at IDs 16,384+ (after
   the original vocab). The first 16,384 IDs — including all 255 special/control tokens — are untouched:

  ┌───────────────────────────────────┬───────────────────┬────────┐
  │               Token               │        ID         │ Status │
  ├───────────────────────────────────┼───────────────────┼────────┤
  │ <unk>                             │ 0                 │ Same   │
  ├───────────────────────────────────┼───────────────────┼────────┤
  │ `<                                │ nospeech          │ >`     │
  ├───────────────────────────────────┼───────────────────┼────────┤
  │ <pad>                             │ 2                 │ Same   │
  ├───────────────────────────────────┼───────────────────┼────────┤
  │ `<                                │ endoftext         │ >`     │
  ├───────────────────────────────────┼───────────────────┼────────┤
  │ `<                                │ startoftranscript │ >`     │
  ├───────────────────────────────────┼───────────────────┼────────┤
  │ `<                                │ pnc               │ >`     │
  ├───────────────────────────────────┼───────────────────┼────────┤
  │ `<                                │ hi                │ >`     │
  ├───────────────────────────────────┼───────────────────┼────────┤
  │ `<                                │ te                │ >`     │
  ├───────────────────────────────────┼───────────────────┼────────┤
  │ All 251 additional special tokens │ 1-254             │ Same   │
  └───────────────────────────────────┴───────────────────┴────────┘

  The special_tokens_map.json from the original model is correct for the extended tokenizer — no changes needed. The inference
  agent can use it as-is.


The user doesn't want to proceed with this tool use. The tool use was rejected (eg. if it was a file edit, the new_string was NOT written to the file). STOP what you are doing and wait for the user to tell you how to proceed.


[Request interrupted by user for tool use]


Good, confirmed on both fronts. Let me run the benchmark now.


python3 benchmark_cohere_transcribe.py \
    --checkpoint /home/ubuntu/training/checkpoints/cohere-transcribe-ckpt-10000 \
    --checkpoint-name ckpt-10000 \
    --languages hi \
    --batch-size 1 2>&1 | tail -30

Smoke test Hindi with device fix


can you run with more BS.


Let me check GPU memory to pick an appropriate batch size.


NVIDIA A100-SXM4-80GB, 81920 MiB, 81156 MiB


A100 80GB with ~81GB free. The model is ~4GB in bf16. Let me start with batch_size=16 for the smoke test on Hindi.


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

Smoke test Hindi BS=16


BS=16 barely using the compute.


Command was manually backgrounded by user with ID: bu742t3nd. Output is being written to: /tmp/claude-1001/-home-ubuntu-training/ed45c711-ad1f-4a43-a0e4-94f0e1c224d8/tasks/bu742t3nd.output


You're right — looking at my inference code, it's actually processing one sample at a time inside the batch loop. Let me rewrite it for proper batched inference.


295	        "formatting_impact": "low",
296	        "worst_languages": [],
297	        "best_languages": [],
298	    }
299
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 inference sample-by-sample (batched mel extraction, sequential generation)."""
307	    results = []
308	    total = len(dataset)
309
310	    for start in range(0, total, batch_size):
311	        end = min(start + batch_size, total)
312	        batch = dataset.select(range(start, end))
313
314	        for i, sample in enumerate(batch):
315	            lang_code = sample["lang_code"]
316	            audio = sample["audio"]
317	            wav = np.array(audio["array"], dtype=np.float32)
318	            sr = audio["sampling_rate"]
319
320	            # Resample if needed
321	            if sr != 16000:
322	                wav_t = torch.from_numpy(wav).unsqueeze(0)
323	                wav_t = torchaudio.functional.resample(wav_t, sr, 16000)
324	                wav = wav_t.squeeze(0).numpy()
325
326	            # Extract mel features
327	            features = processor.feature_extractor(
328	                wav, sampling_rate=16000, return_tensors="pt"
329	            )
330	            input_features = features["input_features"].to(device, dtype=torch.bfloat16)
331	            length = features["length"].to(device)
332
333	            # Build decoder prompt
334	            prompt_str = model.build_prompt(language=lang_code, punctuation=True)
335	            prompt_ids = tokenizer.encode(prompt_str, add_special_tokens=False)
336	            decoder_input_ids = torch.tensor([prompt_ids], device=device)
337
338	            # Generate
339	            with torch.no_grad():
340	                try:
341	                    outputs = model.generate(
342	                        input_features=input_features,
343	                        length=length,
344	                        decoder_input_ids=decoder_input_ids,
345	                        max_new_tokens=max_new_tokens,
346	                    )
347	                    hyp = decode_tokens(tokenizer, outputs[0]).strip()
348	                except Exception as e:
349	                    print(f"  Sample {sample['id']} failed: {e}")
350	                    hyp = ""
351
352	            results.append({
353	                "id": sample["id"],
354	                "language": sample["language"],
355	                "lang_code": lang_code,
356	                "source": sample["source"],
357	                "duration": sample["duration"],
358	                "reference": sample["reference"],
359	                "hypothesis": hyp,

        "formatting_impact": "low",
        "worst_languages": [],
        "best_languages": [],
    }

    return analysis

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

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

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

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

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

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

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

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

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


Now let me check the model's `generate` signature to see if it supports batched inputs.


grep -n "def generate" /home/ubuntu/training/cohere-transcribe-indic/tokenizer_extension/extended_model/modeling_cohere_asr.py


885:    def generate(self, input_features=None, input_ids=None, length=None, attention_mask=None, **kwargs):


885	    def generate(self, input_features=None, input_ids=None, length=None, attention_mask=None, **kwargs):
886	        # If input_ids is provided, use it as decoder_input_ids
887	        # This matches the multimodal encoder-decoder expectation where the prompt is the decoder start
888	        decoder_input_ids = kwargs.pop("decoder_input_ids", None)
889	        if input_ids is not None and decoder_input_ids is None:
890	            decoder_input_ids = input_ids
891	            # We must provide some input_ids to super().generate to avoid validation errors,
892	            # but for encoder-decoder it usually expects encoder input_ids.
893	            # Here input_features is the encoder input.
894	            input_ids = None
895
896	        decoder_attention_mask = kwargs.pop("decoder_attention_mask", None)
897	        if decoder_input_ids is not None and decoder_attention_mask is None:
898	            decoder_attention_mask = torch.ones_like(
899	                decoder_input_ids, dtype=torch.long, device=decoder_input_ids.device
900	            )
901
902	        generation_kwargs = dict(kwargs)
903	        generation_kwargs["input_features"] = input_features
904	        generation_kwargs["length"] = length
905	        generation_kwargs["decoder_input_ids"] = decoder_input_ids
906	        generation_kwargs["decoder_attention_mask"] = decoder_attention_mask
907
908	        decoder_start_token_id = getattr(self.config, "decoder_start_token_id", None)
909	        eos_token_id = getattr(self.config, "eos_token_id", None)
910	        pad_token_id = getattr(self.config, "pad_token_id", None)
911	        if decoder_start_token_id is not None:
912	            generation_kwargs["bos_token_id"] = decoder_start_token_id
913	        if eos_token_id is not None:
914	            generation_kwargs["eos_token_id"] = eos_token_id
915	        if pad_token_id is not None:
916	            generation_kwargs["pad_token_id"] = pad_token_id
917	        if input_ids is not None:
918	            generation_kwargs["input_ids"] = input_ids
919	        if attention_mask is not None:
920	            generation_kwargs["attention_mask"] = attention_mask
921	        if "cache_implementation" not in generation_kwargs:
922	            generation_kwargs["cache_implementation"] = "static"
923
924	        # Fall back to dynamic cache when static cache is incompatible:
925	        # - transformers 4.52-4.55: _supports_static_cache gate + StaticCache
926	        #   reads config.hidden_size which our nested config doesn't expose.
927	        # - transformers >= 5.3: StaticCache.update() API changed (cache_position
928	        #   shape must match key_states, breaking our usage).
929	        if generation_kwargs.get("cache_implementation") == "static":
930	            _skip_static = hasattr(PreTrainedModel, "_supports_static_cache")
931	            if not _skip_static:
932	                import transformers
933
934	                _v = tuple(int(x) for x in transformers.__version__.split(".")[:2])
935	                _skip_static = _v >= (5, 3)
936	            if _skip_static:
937	                generation_kwargs.pop("cache_implementation", None)
938
939	        # We disable_compile for generate() because when passing "cache_implementation"="static"
940	        # transformers will auto-compile the forward pass setting dynamic=False.
941	        # We need dynamic=True to avoid excessive recompilation. Note that this doesn't
942	        # control whether we compile the encoder layers which is set according to
943	        # the transcribe(...,compile=True) flag.
944	        generation_kwargs["disable_compile"] = True
945
946	        return super().generate(**generation_kwargs)
947
948	    def _setup_compile(self, processor=None):
949	        if getattr(self, "_compiled", False):
950	            return
951	        if not hasattr(torch, "compile"):
952	            self._compiled = True
953	            return
954
955	        # Dynamo guards on submodule identity per layer, so each ConformerLayer
956	        # causes a recompilation. Raise the limit so no layers fall back to eager.
957	        needed = len(self.encoder.layers) + 4
958	        if torch._dynamo.config.cache_size_limit < needed:
959	            torch._dynamo.config.cache_size_limit = needed
960
961	        for layer in self.encoder.layers:
962	            layer.forward = torch.compile(layer.forward, dynamic=True)
963
964	        if (

    def generate(self, input_features=None, input_ids=None, length=None, attention_mask=None, **kwargs):
        # If input_ids is provided, use it as decoder_input_ids
        # This matches the multimodal encoder-decoder expectation where the prompt is the decoder start
        decoder_input_ids = kwargs.pop("decoder_input_ids", None)
        if input_ids is not None and decoder_input_ids is None:
            decoder_input_ids = input_ids
            # We must provide some input_ids to super().generate to avoid validation errors,
            # but for encoder-decoder it usually expects encoder input_ids.
            # Here input_features is the encoder input.
            input_ids = None

        decoder_attention_mask = kwargs.pop("decoder_attention_mask", None)
        if decoder_input_ids is not None and decoder_attention_mask is None:
            decoder_attention_mask = torch.ones_like(
                decoder_input_ids, dtype=torch.long, device=decoder_input_ids.device
            )

        generation_kwargs = dict(kwargs)
        generation_kwargs["input_features"] = input_features
        generation_kwargs["length"] = length
        generation_kwargs["decoder_input_ids"] = decoder_input_ids
        generation_kwargs["decoder_attention_mask"] = decoder_attention_mask

        decoder_start_token_id = getattr(self.config, "decoder_start_token_id", None)
        eos_token_id = getattr(self.config, "eos_token_id", None)
        pad_token_id = getattr(self.config, "pad_token_id", None)
        if decoder_start_token_id is not None:
            generation_kwargs["bos_token_id"] = decoder_start_token_id
        if eos_token_id is not None:
            generation_kwargs["eos_token_id"] = eos_token_id
        if pad_token_id is not None:
            generation_kwargs["pad_token_id"] = pad_token_id
        if input_ids is not None:
            generation_kwargs["input_ids"] = input_ids
        if attention_mask is not None:
            generation_kwargs["attention_mask"] = attention_mask
        if "cache_implementation" not in generation_kwargs:
            generation_kwargs["cache_implementation"] = "static"

        # Fall back to dynamic cache when static cache is incompatible:
        # - transformers 4.52-4.55: _supports_static_cache gate + StaticCache
        #   reads config.hidden_size which our nested config doesn't expose.
        # - transformers >= 5.3: StaticCache.update() API changed (cache_position
        #   shape must match key_states, breaking our usage).
        if generation_kwargs.get("cache_implementation") == "static":
            _skip_static = hasattr(PreTrainedModel, "_supports_static_cache")
            if not _skip_static:
                import transformers

                _v = tuple(int(x) for x in transformers.__version__.split(".")[:2])
                _skip_static = _v >= (5, 3)
            if _skip_static:
                generation_kwargs.pop("cache_implementation", None)

        # We disable_compile for generate() because when passing "cache_implementation"="static"
        # transformers will auto-compile the forward pass setting dynamic=False.
        # We need dynamic=True to avoid excessive recompilation. Note that this doesn't
        # control whether we compile the encoder layers which is set according to
        # the transcribe(...,compile=True) flag.
        generation_kwargs["disable_compile"] = True

        return super().generate(**generation_kwargs)

    def _setup_compile(self, processor=None):
        if getattr(self, "_compiled", False):
            return
        if not hasattr(torch, "compile"):
            self._compiled = True
            return

        # Dynamo guards on submodule identity per layer, so each ConformerLayer
        # causes a recompilation. Raise the limit so no layers fall back to eager.
        needed = len(self.encoder.layers) + 4
        if torch._dynamo.config.cache_size_limit < needed:
            torch._dynamo.config.cache_size_limit = needed

        for layer in self.encoder.layers:
            layer.forward = torch.compile(layer.forward, dynamic=True)

        if (


Good - `generate` supports batched `input_features` (it's standard HF generate). The feature extractor's `__call__` already handles lists of audio arrays and pads them. I need to batch the audio and decoder prompts.


Edit

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

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

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

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

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

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

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

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

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

    return results

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