File created successfully at: /home/ubuntu/training/benchmark_cohere_transcribe.py

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

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

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

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

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

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

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

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

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

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

LATIN_SCRIPT_LANGS = {"english"}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    return (len(touched), total_words)

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

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

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

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

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

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

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

        wer_scores = []

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

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

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

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

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

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

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

    return analysis

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

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

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

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

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

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

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

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

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

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

    return results

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

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

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

    # Load model
    from transformers import AutoModelForSpeechSeq2Seq, AutoFeatureExtractor

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    error_analysis = compute_error_analysis(samples_for_errors)

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

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

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

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

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

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

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

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

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

if __name__ == "__main__":
    main()


Updated task #3 status


Now let me check on the checkpoint download.


Bash

cat /tmp/claude-1001/-home-ubuntu-training/ed45c711-ad1f-4a43-a0e4-94f0e1c224d8/tasks/b24ms6xfr.output 2>/dev/null | tail -5


download: s3://ptcheckpoints/cohere-transcribe/04-05-2026/ckpt-10000/processor/tokenizer_config.json to checkpoints/cohere-transcribe-ckpt-10000/tokenizer_config.json
download: s3://ptcheckpoints/cohere-transcribe/04-05-2026/ckpt-10000/processor/processor_config.json to checkpoints/cohere-transcribe-ckpt-10000/processor_config.json
download: s3://ptcheckpoints/cohere-transcribe/04-05-2026/ckpt-10000/processor/processing_cohere_asr.py to checkpoints/cohere-transcribe-ckpt-10000/processing_cohere_asr.py
download: s3://ptcheckpoints/cohere-transcribe/04-05-2026/ckpt-10000/processor/tokenizer.json to checkpoints/cohere-transcribe-ckpt-10000/tokenizer.json
PROCESSOR DONE


total 4.1G
-rw-rw-r-- 1 ubuntu ubuntu 4.2K Apr  5 18:57 config.json
-rw-rw-r-- 1 ubuntu ubuntu 1.7K Apr  5 18:57 configuration_cohere_asr.py
-rw-rw-r-- 1 ubuntu ubuntu  229 Apr  5 18:57 generation_config.json
-rw-rw-r-- 1 ubuntu ubuntu 4.1G Apr  5 19:01 model.safetensors
-rw-rw-r-- 1 ubuntu ubuntu  21K Apr  5 18:57 processing_cohere_asr.py
-rw-rw-r-- 1 ubuntu ubuntu 1.4K Apr  5 18:57 processor_config.json
-rw-rw-r-- 1 ubuntu ubuntu 4.3K Apr  5 18:57 tokenizer_config.json
-rw-rw-r-- 1 ubuntu ubuntu  11M Apr  5 18:57 tokenizer.json


Checkpoint is downloaded. But I notice the checkpoint doesn't have `special_tokens_map.json` which `load_extended_tokenizer` needs, and it doesn't have `modeling_cohere_asr.py` which the model needs for `trust_remote_code=True`. Let me check what files are needed.


# Check what's in config.json for auto_map
python3 -c "import json; c=json.load(open('/home/ubuntu/training/checkpoints/cohere-transcribe-ckpt-10000/config.json')); print(json.dumps(c.get('auto_map', {}), indent=2))"


# Check what's needed: special_tokens_map.json
grep "special_tokens_map" /home/ubuntu/training/cohere-transcribe-indic/tokenizer_utils.py


    special_tokens_file = model_dir / "special_tokens_map.json"
        raise FileNotFoundError(f"Missing special_tokens_map.json in {model_dir}")


The checkpoint references `modeling_cohere_asr.py` and `tokenization_cohere_asr.py` from the auto_map, but those aren't in the R2 checkpoint. They should come from the base model's code. Let me copy the missing files from the extended_model directory.


# Copy missing files from the extended_model (which has the full set of custom code)
CKPT=/home/ubuntu/training/checkpoints/cohere-transcribe-ckpt-10000
SRC=/home/ubuntu/training/cohere-transcribe-indic/tokenizer_extension/extended_model

# Copy modeling code, tokenization code, special_tokens_map, preprocessor_config, vocab_info
for f in modeling_cohere_asr.py tokenization_cohere_asr.py special_tokens_map.json preprocessor_config.json vocab_info.json; do
    if [ ! -f "$CKPT/$f" ] && [ -f "$SRC/$f" ]; then
        cp "$SRC/$f" "$CKPT/$f"
        echo "Copied $f"
    else
        echo "Skipped $f (exists: $([ -f "$CKPT/$f" ] && echo yes || echo no), source: $([ -f "$SRC/$f" ] && echo yes || echo no))"
    fi
done
ls "$CKPT/"


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


Updated task #2 status


Updated task #4 status


Now let me run the benchmark. First a quick smoke test on a couple samples.


python benchmark_cohere_transcribe.py \
    --checkpoint /home/ubuntu/training/checkpoints/cohere-transcribe-ckpt-10000 \
    --checkpoint-name ckpt-10000 \
    --languages hi en \
    --batch-size 1 2>&1 | head -40

Quick smoke test on 2 languages


/bin/bash: line 5: python: command not found


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

Quick smoke test on hi+en languages


<persisted-output>
Output too large (34.2KB). Full output saved to: /home/ubuntu/.claude/projects/-home-ubuntu-training/ed45c711-ad1f-4a43-a0e4-94f0e1c224d8/tool-results/b5i371ozs.txt

Preview (first 2KB):
Skipping import of cpp extensions due to incompatible torch version 2.7.0+cu126 for torchao version 0.16.0             Please see https://github.com/pytorch/ao/issues/2919 for more info
`torch_dtype` is deprecated! Use `dtype` instead!
Some weights of the model checkpoint at /home/ubuntu/training/checkpoints/cohere-transcribe-ckpt-10000 were not used when initializing CohereAsrForConditionalGeneration: ['encoder_encoder_decoder_proj.bias', 'encoder_encoder_decoder_proj.weight', 'transf_decoder._transf_decoder._decoder.layers.0.first_sub_layer.key_net.bias', 'transf_decoder._transf_decoder._decoder.layers.0.first_sub_layer.key_net.weight', 'transf_decoder._transf_decoder._decoder.layers.0.first_sub_layer.out_projection.bias', 'transf_decoder._transf_decoder._decoder.layers.0.first_sub_layer.out_projection.weight', 'transf_decoder._transf_decoder._decoder.layers.0.first_sub_layer.query_net.bias', 'transf_decoder._transf_decoder._decoder.layers.0.first_sub_layer.query_net.weight', 'transf_decoder._transf_decoder._decoder.layers.0.first_sub_layer.value_net.bias', 'transf_decoder._transf_decoder._decoder.layers.0.first_sub_layer.value_net.weight', 'transf_decoder._transf_decoder._decoder.layers.0.layer_norm_1.bias', 'transf_decoder._transf_decoder._decoder.layers.0.layer_norm_1.weight', 'transf_decoder._transf_decoder._decoder.layers.0.layer_norm_2.bias', 'transf_decoder._transf_decoder._decoder.layers.0.layer_norm_2.weight', 'transf_decoder._transf_decoder._decoder.layers.0.layer_norm_3.bias', 'transf_decoder._transf_decoder._decoder.layers.0.layer_norm_3.weight', 'transf_decoder._transf_decoder._decoder.layers.0.second_sub_layer.key_net.bias', 'transf_decoder._transf_decoder._decoder.layers.0.second_sub_layer.key_net.weight', 'transf_decoder._transf_decoder._decoder.layers.0.second_sub_layer.out_projection.bias', 'transf_decoder._transf_decoder._decoder.layers.0.second_sub_layer.out_projection.weight', 'transf_decoder._transf_decoder._decoder.layers.0.second_sub_layer.
...
</persisted-output>

Skipping import of cpp extensions due to incompatible torch version 2.7.0+cu126 for torchao version 0.16.0             Please see https://github.com/pytorch/ao/issues/2919 for more info
`torch_dtype` is deprecated! Use `dtype` instead!
Some weights of the model checkpoint at /home/ubuntu/training/checkpoints/cohere-transcribe-ckpt-10000 were not used when initializing CohereAsrForConditionalGeneration: ['encoder_encoder_decoder_proj.bias', 'encoder_encoder_decoder_proj.weight', 'transf_decoder._transf_decoder._decoder.layers.0.first_sub_layer.key_net.bias', 'transf_decoder._transf_decoder._decoder.layers.0.first_sub_layer.key_net.weight', 'transf_decoder._transf_decoder._decoder.layers.0.first_sub_layer.out_projection.bias', 'transf_decoder._transf_decoder._decoder.layers.0.first_sub_layer.out_projection.weight', 'transf_decoder._transf_decoder._decoder.layers.0.first_sub_layer.query_net.bias', 'transf_decoder._transf_decoder._decoder.layers.0.first_sub_layer.query_net.weight', 'transf_decoder._transf_decoder._decoder.layers.0.first_sub_layer.value_net.bias', 'transf_decoder._transf_decoder._decoder.layers.0.first_sub_layer.value_net.weight', 'transf_decoder._transf_decoder._decoder.layers.0.layer_norm_1.bias', 'transf_decoder._transf_decoder._decoder.layers.0.layer_norm_1.weight', 'transf_decoder._transf_decoder._decoder.layers.0.layer_norm_2.bias', 'transf_decoder._transf_decoder._decoder.layers.0.layer_norm_2.weight', 'transf_decoder._transf_decoder._decoder.layers.0.layer_norm_3.bias', 'transf_decoder._transf_decoder._decoder.layers.0.layer_norm_3.weight', 'transf_decoder._transf_decoder._decoder.layers.0.second_sub_layer.key_net.bias', 'transf_decoder._transf_decoder._decoder.layers.0.second_sub_layer.key_net.weight', 'transf_decoder._transf_decoder._decoder.layers.0.second_sub_layer.out_projection.bias', 'transf_decoder._transf_decoder._decoder.layers.0.second_sub_layer.out_projection.weight', 'transf_decoder._transf_decoder._decoder.layers.0.second_sub_layer.query_net.bias', 'transf_decoder._transf_decoder._decoder.layers.0.second_sub_layer.query_net.weight', 'transf_decoder._transf_decoder._decoder.layers.0.second_sub_layer.value_net.bias', 'transf_decoder._transf_decoder._decoder.layers.0.second_sub_layer.value_net.weight', 'transf_decoder._transf_decoder._decoder.layers.0.third_sub_layer.dense_in.bias', 'transf_decoder._transf_decoder._decoder.layers.0.third_sub_layer.dense_in.weight', 'transf_decoder._transf_decoder._decoder.layers.0.third_sub_layer.dense_out.bias', 'transf_decoder._transf_decoder._decoder.layers.0.third_sub_layer.dense_out.weight', 'transf_decoder._transf_decoder._decoder.layers.1.first_sub_layer.key_net.bias', 'transf_decoder._transf_decoder._decoder.layers.1.first_sub_layer.key_net.weight', 'transf_decoder._transf_decoder._decoder.layers.1.first_sub_layer.out_projection.bias', 'transf_decoder._transf_decoder._decoder.layers.1.first_sub_layer.out_projection.weight', 'transf_decoder._transf_decoder._decoder.layers.1.first_sub_layer.query_net.bias', 'transf_decoder._transf_decoder._decoder.layers.1.first_sub_layer.query_net.weight', 'transf_decoder._transf_decoder._decoder.layers.1.first_sub_layer.value_net.bias', 'transf_decoder._transf_decoder._decoder.layers.1.first_sub_layer.value_net.weight', 'transf_decoder._transf_decoder._decoder.layers.1.layer_norm_1.bias', 'transf_decoder._transf_decoder._decoder.layers.1.layer_norm_1.weight', 'transf_decoder._transf_decoder._decoder.layers.1.layer_norm_2.bias', 'transf_decoder._transf_decoder._decoder.layers.1.layer_norm_2.weight', 'transf_decoder._transf_decoder._decoder.layers.1.layer_norm_3.bias', 'transf_decoder._transf_decoder._decoder.layers.1.layer_norm_3.weight', 'transf_decoder._transf_decoder._decoder.layers.1.second_sub_layer.key_net.bias', 'transf_decoder._transf_decoder._decoder.layers.1.second_sub_layer.key_net.weight', 'transf_decoder._transf_decoder._decoder.layers.1.second_sub_layer.out_projection.bias', 'transf_decoder._transf_decoder._decoder.layers.1.second_sub_layer.out_projection.weight', 'transf_decoder._transf_decoder._decoder.layers.1.second_sub_layer.query_net.bias', 'transf_decoder._transf_decoder._decoder.layers.1.second_sub_layer.query_net.weight', 'transf_decoder._transf_decoder._decoder.layers.1.second_sub_layer.value_net.bias', 'transf_decoder._transf_decoder._decoder.layers.1.second_sub_layer.value_net.weight', 'transf_decoder._transf_decoder._decoder.layers.1.third_sub_layer.dense_in.bias', 'transf_decoder._transf_decoder._decoder.layers.1.third_sub_layer.dense_in.weight', 'transf_decoder._transf_decoder._decoder.layers.1.third_sub_layer.dense_out.bias', 'transf_decoder._transf_decoder._decoder.layers.1.third_sub_layer.dense_out.weight', 'transf_decoder._transf_decoder._decoder.layers.2.first_sub_layer.key_net.bias', 'transf_decoder._transf_decoder._decoder.layers.2.first_sub_layer.key_net.weight', 'transf_decoder._transf_decoder._decoder.layers.2.first_sub_layer.out_projection.bias', 'transf_decoder._transf_decoder._decoder.layers.2.first_sub_layer.out_projection.weight', 'transf_decoder._transf_decoder._decoder.layers.2.first_sub_layer.query_net.bias', 'transf_decoder._transf_decoder._decoder.layers.2.first_sub_layer.query_net.weight', 'transf_decoder._transf_decoder._decoder.layers.2.first_sub_layer.value_net.bias', 'transf_decoder._transf_decoder._decoder.layers.2.first_sub_layer.value_net.weight', 'transf_decoder._transf_decoder._decoder.layers.2.layer_norm_1.bias', 'transf_decoder._transf_decoder._decoder.layers.2.layer_norm_1.weight', 'transf_decoder._transf_decoder._decoder.layers.2.layer_norm_2.bias', 'transf_decoder._transf_decoder._decoder.layers.2.layer_norm_2.weight', 'transf_decoder._transf_decoder._decoder.layers.2.layer_norm_3.bias', 'transf_decoder._transf_decoder._decoder.layers.2.layer_norm_3.weight', 'transf_decoder._transf_decoder._decoder.layers.2.second_sub_layer.key_net.bias', 'transf_decoder._transf_decoder._decoder.layers.2.second_sub_layer.key_net.weight', 'transf_decoder._transf_decoder._decoder.layers.2.second_sub_layer.out_projection.bias', 'transf_decoder._transf_decoder._decoder.layers.2.second_sub_layer.out_projection.weight', 'transf_decoder._transf_decoder._decoder.layers.2.second_sub_layer.query_net.bias', 'transf_decoder._transf_decoder._decoder.layers.2.second_sub_layer.query_net.weight', 'transf_decoder._transf_decoder._decoder.layers.2.second_sub_layer.value_net.bias', 'transf_decoder._transf_decoder._decoder.layers.2.second_sub_layer.value_net.weight', 'transf_decoder._transf_decoder._decoder.layers.2.third_sub_layer.dense_in.bias', 'transf_decoder._transf_decoder._decoder.layers.2.third_sub_layer.dense_in.weight', 'transf_decoder._transf_decoder._decoder.layers.2.third_sub_layer.dense_out.bias', 'transf_decoder._transf_decoder._decoder.layers.2.third_sub_layer.dense_out.weight', 'transf_decoder._transf_decoder._decoder.layers.3.first_sub_layer.key_net.bias', 'transf_decoder._transf_decoder._decoder.layers.3.first_sub_layer.key_net.weight', 'transf_decoder._transf_decoder._decoder.layers.3.first_sub_layer.out_projection.bias', 'transf_decoder._transf_decoder._decoder.layers.3.first_sub_layer.out_projection.weight', 'transf_decoder._transf_decoder._decoder.layers.3.first_sub_layer.query_net.bias', 'transf_decoder._transf_decoder._decoder.layers.3.first_sub_layer.query_net.weight', 'transf_decoder._transf_decoder._decoder.layers.3.first_sub_layer.value_net.bias', 'transf_decoder._transf_decoder._decoder.layers.3.first_sub_layer.value_net.weight', 'transf_decoder._transf_decoder._decoder.layers.3.layer_norm_1.bias', 'transf_decoder._transf_decoder._decoder.layers.3.layer_norm_1.weight', 'transf_decoder._transf_decoder._decoder.layers.3.layer_norm_2.bias', 'transf_decoder._transf_decoder._decoder.layers.3.layer_norm_2.weight', 'transf_decoder._transf_decoder._decoder.layers.3.layer_norm_3.bias', 'transf_decoder._transf_decoder._decoder.layers.3.layer_norm_3.weight', 'transf_decoder._transf_decoder._decoder.layers.3.second_sub_layer.key_net.bias', 'transf_decoder._transf_decoder._decoder.layers.3.second_sub_layer.key_net.weight', 'transf_decoder._transf_decoder._decoder.layers.3.second_sub_layer.out_projection.bias', 'transf_decoder._transf_decoder._decoder.layers.3.second_sub_layer.out_projection.weight', 'transf_decoder._transf_decoder._decoder.layers.3.second_sub_layer.query_net.bias', 'transf_decoder._transf_decoder._decoder.layers.3.second_sub_layer.query_net.weight', 'transf_decoder._transf_decoder._decoder.layers.3.second_sub_layer.value_net.bias', 'transf_decoder._transf_decoder._decoder.layers.3.second_sub_layer.value_net.weight', 'transf_decoder._transf_decoder._decoder.layers.3.third_sub_layer.dense_in.bias', 'transf_decoder._transf_decoder._decoder.layers.3.third_sub_layer.dense_in.weight', 'transf_decoder._transf_decoder._decoder.layers.3.third_sub_layer.dense_out.bias', 'transf_decoder._transf_decoder._decoder.layers.3.third_sub_layer.dense_out.weight', 'transf_decoder._transf_decoder._decoder.layers.4.first_sub_layer.key_net.bias', 'transf_decoder._transf_decoder._decoder.layers.4.first_sub_layer.key_net.weight', 'transf_decoder._transf_decoder._decoder.layers.4.first_sub_layer.out_projection.bias', 'transf_decoder._transf_decoder._decoder.layers.4.first_sub_layer.out_projection.weight', 'transf_decoder._transf_decoder._decoder.layers.4.first_sub_layer.query_net.bias', 'transf_decoder._transf_decoder._decoder.layers.4.first_sub_layer.query_net.weight', 'transf_decoder._transf_decoder._decoder.layers.4.first_sub_layer.value_net.bias', 'transf_decoder._transf_decoder._decoder.layers.4.first_sub_layer.value_net.weight', 'transf_decoder._transf_decoder._decoder.layers.4.layer_norm_1.bias', 'transf_decoder._transf_decoder._decoder.layers.4.layer_norm_1.weight', 'transf_decoder._transf_decoder._decoder.layers.4.layer_norm_2.bias', 'transf_decoder._transf_decoder._decoder.layers.4.layer_norm_2.weight', 'transf_decoder._transf_decoder._decoder.layers.4.layer_norm_3.bias', 'transf_decoder._transf_decoder._decoder.layers.4.layer_norm_3.weight', 'transf_decoder._transf_decoder._decoder.layers.4.second_sub_layer.key_net.bias', 'transf_decoder._transf_decoder._decoder.layers.4.second_sub_layer.key_net.weight', 'transf_decoder._transf_decoder._decoder.layers.4.second_sub_layer.out_projection.bias', 'transf_decoder._transf_decoder._decoder.layers.4.second_sub_layer.out_projection.weight', 'transf_decoder._transf_decoder._decoder.layers.4.second_sub_layer.query_net.bias', 'transf_decoder._transf_decoder._decoder.layers.4.second_sub_layer.query_net.weight', 'transf_decoder._transf_decoder._decoder.layers.4.second_sub_layer.value_net.bias', 'transf_decoder._transf_decoder._decoder.layers.4.second_sub_layer.value_net.weight', 'transf_decoder._transf_decoder._decoder.layers.4.third_sub_layer.dense_in.bias', 'transf_decoder._transf_decoder._decoder.layers.4.third_sub_layer.dense_in.weight', 'transf_decoder._transf_decoder._decoder.layers.4.third_sub_layer.dense_out.bias', 'transf_decoder._transf_decoder._decoder.layers.4.third_sub_layer.dense_out.weight', 'transf_decoder._transf_decoder._decoder.layers.5.first_sub_layer.key_net.bias', 'transf_decoder._transf_decoder._decoder.layers.5.first_sub_layer.key_net.weight', 'transf_decoder._transf_decoder._decoder.layers.5.first_sub_layer.out_projection.bias', 'transf_decoder._transf_decoder._decoder.layers.5.first_sub_layer.out_projection.weight', 'transf_decoder._transf_decoder._decoder.layers.5.first_sub_layer.query_net.bias', 'transf_decoder._transf_decoder._decoder.layers.5.first_sub_layer.query_net.weight', 'transf_decoder._transf_decoder._decoder.layers.5.first_sub_layer.value_net.bias', 'transf_decoder._transf_decoder._decoder.layers.5.first_sub_layer.value_net.weight', 'transf_decoder._transf_decoder._decoder.layers.5.layer_norm_1.bias', 'transf_decoder._transf_decoder._decoder.layers.5.layer_norm_1.weight', 'transf_decoder._transf_decoder._decoder.layers.5.layer_norm_2.bias', 'transf_decoder._transf_decoder._decoder.layers.5.layer_norm_2.weight', 'transf_decoder._transf_decoder._decoder.layers.5.layer_norm_3.bias', 'transf_decoder._transf_decoder._decoder.layers.5.layer_norm_3.weight', 'transf_decoder._transf_decoder._decoder.layers.5.second_sub_layer.key_net.bias', 'transf_decoder._transf_decoder._decoder.layers.5.second_sub_layer.key_net.weight', 'transf_decoder._transf_decoder._decoder.layers.5.second_sub_layer.out_projection.bias', 'transf_decoder._transf_decoder._decoder.layers.5.second_sub_layer.out_projection.weight', 'transf_decoder._transf_decoder._decoder.layers.5.second_sub_layer.query_net.bias', 'transf_decoder._transf_decoder._decoder.layers.5.second_sub_layer.query_net.weight', 'transf_decoder._transf_decoder._decoder.layers.5.second_sub_layer.value_net.bias', 'transf_decoder._transf_decoder._decoder.layers.5.second_sub_layer.value_net.weight', 'transf_decoder._transf_decoder._decoder.layers.5.third_sub_layer.dense_in.bias', 'transf_decoder._transf_decoder._decoder.layers.5.third_sub_layer.dense_in.weight', 'transf_decoder._transf_decoder._decoder.layers.5.third_sub_layer.dense_out.bias', 'transf_decoder._transf_decoder._decoder.layers.5.third_sub_layer.dense_out.weight', 'transf_decoder._transf_decoder._decoder.layers.6.first_sub_layer.key_net.bias', 'transf_decoder._transf_decoder._decoder.layers.6.first_sub_layer.key_net.weight', 'transf_decoder._transf_decoder._decoder.layers.6.first_sub_layer.out_projection.bias', 'transf_decoder._transf_decoder._decoder.layers.6.first_sub_layer.out_projection.weight', 'transf_decoder._transf_decoder._decoder.layers.6.first_sub_layer.query_net.bias', 'transf_decoder._transf_decoder._decoder.layers.6.first_sub_layer.query_net.weight', 'transf_decoder._transf_decoder._decoder.layers.6.first_sub_layer.value_net.bias', 'transf_decoder._transf_decoder._decoder.layers.6.first_sub_layer.value_net.weight', 'transf_decoder._transf_decoder._decoder.layers.6.layer_norm_1.bias', 'transf_decoder._transf_decoder._decoder.layers.6.layer_norm_1.weight', 'transf_decoder._transf_decoder._decoder.layers.6.layer_norm_2.bias', 'transf_decoder._transf_decoder._decoder.layers.6.layer_norm_2.weight', 'transf_decoder._transf_decoder._decoder.layers.6.layer_norm_3.bias', 'transf_decoder._transf_decoder._decoder.layers.6.layer_norm_3.weight', 'transf_decoder._transf_decoder._decoder.layers.6.second_sub_layer.key_net.bias', 'transf_decoder._transf_decoder._decoder.layers.6.second_sub_layer.key_net.weight', 'transf_decoder._transf_decoder._decoder.layers.6.second_sub_layer.out_projection.bias', 'transf_decoder._transf_decoder._decoder.layers.6.second_sub_layer.out_projection.weight', 'transf_decoder._transf_decoder._decoder.layers.6.second_sub_layer.query_net.bias', 'transf_decoder._transf_decoder._decoder.layers.6.second_sub_layer.query_net.weight', 'transf_decoder._transf_decoder._decoder.layers.6.second_sub_layer.value_net.bias', 'transf_decoder._transf_decoder._decoder.layers.6.second_sub_layer.value_net.weight', 'transf_decoder._transf_decoder._decoder.layers.6.third_sub_layer.dense_in.bias', 'transf_decoder._transf_decoder._decoder.layers.6.third_sub_layer.dense_in.weight', 'transf_decoder._transf_decoder._decoder.layers.6.third_sub_layer.dense_out.bias', 'transf_decoder._transf_decoder._decoder.layers.6.third_sub_layer.dense_out.weight', 'transf_decoder._transf_decoder._decoder.layers.7.first_sub_layer.key_net.bias', 'transf_decoder._transf_decoder._decoder.layers.7.first_sub_layer.key_net.weight', 'transf_decoder._transf_decoder._decoder.layers.7.first_sub_layer.out_projection.bias', 'transf_decoder._transf_decoder._decoder.layers.7.first_sub_layer.out_projection.weight', 'transf_decoder._transf_decoder._decoder.layers.7.first_sub_layer.query_net.bias', 'transf_decoder._transf_decoder._decoder.layers.7.first_sub_layer.query_net.weight', 'transf_decoder._transf_decoder._decoder.layers.7.first_sub_layer.value_net.bias', 'transf_decoder._transf_decoder._decoder.layers.7.first_sub_layer.value_net.weight', 'transf_decoder._transf_decoder._decoder.layers.7.layer_norm_1.bias', 'transf_decoder._transf_decoder._decoder.layers.7.layer_norm_1.weight', 'transf_decoder._transf_decoder._decoder.layers.7.layer_norm_2.bias', 'transf_decoder._transf_decoder._decoder.layers.7.layer_norm_2.weight', 'transf_decoder._transf_decoder._decoder.layers.7.layer_norm_3.bias', 'transf_decoder._transf_decoder._decoder.layers.7.layer_norm_3.weight', 'transf_decoder._transf_decoder._decoder.layers.7.second_sub_layer.key_net.bias', 'transf_decoder._transf_decoder._decoder.layers.7.second_sub_layer.key_net.weight', 'transf_decoder._transf_decoder._decoder.layers.7.second_sub_layer.out_projection.bias', 'transf_decoder._transf_decoder._decoder.layers.7.second_sub_layer.out_projection.weight', 'transf_decoder._transf_decoder._decoder.layers.7.second_sub_layer.query_net.bias', 'transf_decoder._transf_decoder._decoder.layers.7.second_sub_layer.query_net.weight', 'transf_decoder._transf_decoder._decoder.layers.7.second_sub_layer.value_net.bias', 'transf_decoder._transf_decoder._decoder.layers.7.second_sub_layer.value_net.weight', 'transf_decoder._transf_decoder._decoder.layers.7.third_sub_layer.dense_in.bias', 'transf_decoder._transf_decoder._decoder.layers.7.third_sub_layer.dense_in.weight', 'transf_decoder._transf_decoder._decoder.layers.7.third_sub_layer.dense_out.bias', 'transf_decoder._transf_decoder._decoder.layers.7.third_sub_layer.dense_out.weight']
- This IS expected if you are initializing CohereAsrForConditionalGeneration from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
- This IS NOT expected if you are initializing CohereAsrForConditionalGeneration from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
Some weights of CohereAsrForConditionalGeneration were not initialized from the model checkpoint at /home/ubuntu/training/checkpoints/cohere-transcribe-ckpt-10000 and are newly initialized: ['encoder_decoder_proj.bias', 'encoder_decoder_proj.weight', 'transf_decoder._decoder.layers.0.first_sub_layer.key_net.bias', 'transf_decoder._decoder.layers.0.first_sub_layer.key_net.weight', 'transf_decoder._decoder.layers.0.first_sub_layer.out_projection.bias', 'transf_decoder._decoder.layers.0.first_sub_layer.out_projection.weight', 'transf_decoder._decoder.layers.0.first_sub_layer.query_net.bias', 'transf_decoder._decoder.layers.0.first_sub_layer.query_net.weight', 'transf_decoder._decoder.layers.0.first_sub_layer.value_net.bias', 'transf_decoder._decoder.layers.0.first_sub_layer.value_net.weight', 'transf_decoder._decoder.layers.0.layer_norm_1.bias', 'transf_decoder._decoder.layers.0.layer_norm_1.weight', 'transf_decoder._decoder.layers.0.layer_norm_2.bias', 'transf_decoder._decoder.layers.0.layer_norm_2.weight', 'transf_decoder._decoder.layers.0.layer_norm_3.bias', 'transf_decoder._decoder.layers.0.layer_norm_3.weight', 'transf_decoder._decoder.layers.0.second_sub_layer.key_net.bias', 'transf_decoder._decoder.layers.0.second_sub_layer.key_net.weight', 'transf_decoder._decoder.layers.0.second_sub_layer.out_projection.bias', 'transf_decoder._decoder.layers.0.second_sub_layer.out_projection.weight', 'transf_decoder._decoder.layers.0.second_sub_layer.query_net.bias', 'transf_decoder._decoder.layers.0.second_sub_layer.query_net.weight', 'transf_decoder._decoder.layers.0.second_sub_layer.value_net.bias', 'transf_decoder._decoder.layers.0.second_sub_layer.value_net.weight', 'transf_decoder._decoder.layers.0.third_sub_layer.dense_in.bias', 'transf_decoder._decoder.layers.0.third_sub_layer.dense_in.weight', 'transf_decoder._decoder.layers.0.third_sub_layer.dense_out.bias', 'transf_decoder._decoder.layers.0.third_sub_layer.dense_out.weight', 'transf_decoder._decoder.layers.1.first_sub_layer.key_net.bias', 'transf_decoder._decoder.layers.1.first_sub_layer.key_net.weight', 'transf_decoder._decoder.layers.1.first_sub_layer.out_projection.bias', 'transf_decoder._decoder.layers.1.first_sub_layer.out_projection.weight', 'transf_decoder._decoder.layers.1.first_sub_layer.query_net.bias', 'transf_decoder._decoder.layers.1.first_sub_layer.query_net.weight', 'transf_decoder._decoder.layers.1.first_sub_layer.value_net.bias', 'transf_decoder._decoder.layers.1.first_sub_layer.value_net.weight', 'transf_decoder._decoder.layers.1.layer_norm_1.bias', 'transf_decoder._decoder.layers.1.layer_norm_1.weight', 'transf_decoder._decoder.layers.1.layer_norm_2.bias', 'transf_decoder._decoder.layers.1.layer_norm_2.weight', 'transf_decoder._decoder.layers.1.layer_norm_3.bias', 'transf_decoder._decoder.layers.1.layer_norm_3.weight', 'transf_decoder._decoder.layers.1.second_sub_layer.key_net.bias', 'transf_decoder._decoder.layers.1.second_sub_layer.key_net.weight', 'transf_decoder._decoder.layers.1.second_sub_layer.out_projection.bias', 'transf_decoder._decoder.layers.1.second_sub_layer.out_projection.weight', 'transf_decoder._decoder.layers.1.second_sub_layer.query_net.bias', 'transf_decoder._decoder.layers.1.second_sub_layer.query_net.weight', 'transf_decoder._decoder.layers.1.second_sub_layer.value_net.bias', 'transf_decoder._decoder.layers.1.second_sub_layer.value_net.weight', 'transf_decoder._decoder.layers.1.third_sub_layer.dense_in.bias', 'transf_decoder._decoder.layers.1.third_sub_layer.dense_in.weight', 'transf_decoder._decoder.layers.1.third_sub_layer.dense_out.bias', 'transf_decoder._decoder.layers.1.third_sub_layer.dense_out.weight', 'transf_decoder._decoder.layers.2.first_sub_layer.key_net.bias', 'transf_decoder._decoder.layers.2.first_sub_layer.key_net.weight', 'transf_decoder._decoder.layers.2.first_sub_layer.out_projection.bias', 'transf_decoder._decoder.layers.2.first_sub_layer.out_projection.weight', 'transf_decoder._decoder.layers.2.first_sub_layer.query_net.bias', 'transf_decoder._decoder.layers.2.first_sub_layer.query_net.weight', 'transf_decoder._decoder.layers.2.first_sub_layer.value_net.bias', 'transf_decoder._decoder.layers.2.first_sub_layer.value_net.weight', 'transf_decoder._decoder.layers.2.layer_norm_1.bias', 'transf_decoder._decoder.layers.2.layer_norm_1.weight', 'transf_decoder._decoder.layers.2.layer_norm_2.bias', 'transf_decoder._decoder.layers.2.layer_norm_2.weight', 'transf_decoder._decoder.layers.2.layer_norm_3.bias', 'transf_decoder._decoder.layers.2.layer_norm_3.weight', 'transf_decoder._decoder.layers.2.second_sub_layer.key_net.bias', 'transf_decoder._decoder.layers.2.second_sub_layer.key_net.weight', 'transf_decoder._decoder.layers.2.second_sub_layer.out_projection.bias', 'transf_decoder._decoder.layers.2.second_sub_layer.out_projection.weight', 'transf_decoder._decoder.layers.2.second_sub_layer.query_net.bias', 'transf_decoder._decoder.layers.2.second_sub_layer.query_net.weight', 'transf_decoder._decoder.layers.2.second_sub_layer.value_net.bias', 'transf_decoder._decoder.layers.2.second_sub_layer.value_net.weight', 'transf_decoder._decoder.layers.2.third_sub_layer.dense_in.bias', 'transf_decoder._decoder.layers.2.third_sub_layer.dense_in.weight', 'transf_decoder._decoder.layers.2.third_sub_layer.dense_out.bias', 'transf_decoder._decoder.layers.2.third_sub_layer.dense_out.weight', 'transf_decoder._decoder.layers.3.first_sub_layer.key_net.bias', 'transf_decoder._decoder.layers.3.first_sub_layer.key_net.weight', 'transf_decoder._decoder.layers.3.first_sub_layer.out_projection.bias', 'transf_decoder._decoder.layers.3.first_sub_layer.out_projection.weight', 'transf_decoder._decoder.layers.3.first_sub_layer.query_net.bias', 'transf_decoder._decoder.layers.3.first_sub_layer.query_net.weight', 'transf_decoder._decoder.layers.3.first_sub_layer.value_net.bias', 'transf_decoder._decoder.layers.3.first_sub_layer.value_net.weight', 'transf_decoder._decoder.layers.3.layer_norm_1.bias', 'transf_decoder._decoder.layers.3.layer_norm_1.weight', 'transf_decoder._decoder.layers.3.layer_norm_2.bias', 'transf_decoder._decoder.layers.3.layer_norm_2.weight', 'transf_decoder._decoder.layers.3.layer_norm_3.bias', 'transf_decoder._decoder.layers.3.layer_norm_3.weight', 'transf_decoder._decoder.layers.3.second_sub_layer.key_net.bias', 'transf_decoder._decoder.layers.3.second_sub_layer.key_net.weight', 'transf_decoder._decoder.layers.3.second_sub_layer.out_projection.bias', 'transf_decoder._decoder.layers.3.second_sub_layer.out_projection.weight', 'transf_decoder._decoder.layers.3.second_sub_layer.query_net.bias', 'transf_decoder._decoder.layers.3.second_sub_layer.query_net.weight', 'transf_decoder._decoder.layers.3.second_sub_layer.value_net.bias', 'transf_decoder._decoder.layers.3.second_sub_layer.value_net.weight', 'transf_decoder._decoder.layers.3.third_sub_layer.dense_in.bias', 'transf_decoder._decoder.layers.3.third_sub_layer.dense_in.weight', 'transf_decoder._decoder.layers.3.third_sub_layer.dense_out.bias', 'transf_decoder._decoder.layers.3.third_sub_layer.dense_out.weight', 'transf_decoder._decoder.layers.4.first_sub_layer.key_net.bias', 'transf_decoder._decoder.layers.4.first_sub_layer.key_net.weight', 'transf_decoder._decoder.layers.4.first_sub_layer.out_projection.bias', 'transf_decoder._decoder.layers.4.first_sub_layer.out_projection.weight', 'transf_decoder._decoder.layers.4.first_sub_layer.query_net.bias', 'transf_decoder._decoder.layers.4.first_sub_layer.query_net.weight', 'transf_decoder._decoder.layers.4.first_sub_layer.value_net.bias', 'transf_decoder._decoder.layers.4.first_sub_layer.value_net.weight', 'transf_decoder._decoder.layers.4.layer_norm_1.bias', 'transf_decoder._decoder.layers.4.layer_norm_1.weight', 'transf_decoder._decoder.layers.4.layer_norm_2.bias', 'transf_decoder._decoder.layers.4.layer_norm_2.weight', 'transf_decoder._decoder.layers.4.layer_norm_3.bias', 'transf_decoder._decoder.layers.4.layer_norm_3.weight', 'transf_decoder._decoder.layers.4.second_sub_layer.key_net.bias', 'transf_decoder._decoder.layers.4.second_sub_layer.key_net.weight', 'transf_decoder._decoder.layers.4.second_sub_layer.out_projection.bias', 'transf_decoder._decoder.layers.4.second_sub_layer.out_projection.weight', 'transf_decoder._decoder.layers.4.second_sub_layer.query_net.bias', 'transf_decoder._decoder.layers.4.second_sub_layer.query_net.weight', 'transf_decoder._decoder.layers.4.second_sub_layer.value_net.bias', 'transf_decoder._decoder.layers.4.second_sub_layer.value_net.weight', 'transf_decoder._decoder.layers.4.third_sub_layer.dense_in.bias', 'transf_decoder._decoder.layers.4.third_sub_layer.dense_in.weight', 'transf_decoder._decoder.layers.4.third_sub_layer.dense_out.bias', 'transf_decoder._decoder.layers.4.third_sub_layer.dense_out.weight', 'transf_decoder._decoder.layers.5.first_sub_layer.key_net.bias', 'transf_decoder._decoder.layers.5.first_sub_layer.key_net.weight', 'transf_decoder._decoder.layers.5.first_sub_layer.out_projection.bias', 'transf_decoder._decoder.layers.5.first_sub_layer.out_projection.weight', 'transf_decoder._decoder.layers.5.first_sub_layer.query_net.bias', 'transf_decoder._decoder.layers.5.first_sub_layer.query_net.weight', 'transf_decoder._decoder.layers.5.first_sub_layer.value_net.bias', 'transf_decoder._decoder.layers.5.first_sub_layer.value_net.weight', 'transf_decoder._decoder.layers.5.layer_norm_1.bias', 'transf_decoder._decoder.layers.5.layer_norm_1.weight', 'transf_decoder._decoder.layers.5.layer_norm_2.bias', 'transf_decoder._decoder.layers.5.layer_norm_2.weight', 'transf_decoder._decoder.layers.5.layer_norm_3.bias', 'transf_decoder._decoder.layers.5.layer_norm_3.weight', 'transf_decoder._decoder.layers.5.second_sub_layer.key_net.bias', 'transf_decoder._decoder.layers.5.second_sub_layer.key_net.weight', 'transf_decoder._decoder.layers.5.second_sub_layer.out_projection.bias', 'transf_decoder._decoder.layers.5.second_sub_layer.out_projection.weight', 'transf_decoder._decoder.layers.5.second_sub_layer.query_net.bias', 'transf_decoder._decoder.layers.5.second_sub_layer.query_net.weight', 'transf_decoder._decoder.layers.5.second_sub_layer.value_net.bias', 'transf_decoder._decoder.layers.5.second_sub_layer.value_net.weight', 'transf_decoder._decoder.layers.5.third_sub_layer.dense_in.bias', 'transf_decoder._decoder.layers.5.third_sub_layer.dense_in.weight', 'transf_decoder._decoder.layers.5.third_sub_layer.dense_out.bias', 'transf_decoder._decoder.layers.5.third_sub_layer.dense_out.weight', 'transf_decoder._decoder.layers.6.first_sub_layer.key_net.bias', 'transf_decoder._decoder.layers.6.first_sub_layer.key_net.weight', 'transf_decoder._decoder.layers.6.first_sub_layer.out_projection.bias', 'transf_decoder._decoder.layers.6.first_sub_layer.out_projection.weight', 'transf_decoder._decoder.layers.6.first_sub_layer.query_net.bias', 'transf_decoder._decoder.layers.6.first_sub_layer.query_net.weight', 'transf_decoder._decoder.layers.6.first_sub_layer.value_net.bias', 'transf_decoder._decoder.layers.6.first_sub_layer.value_net.weight', 'transf_decoder._decoder.layers.6.layer_norm_1.bias', 'transf_decoder._decoder.layers.6.layer_norm_1.weight', 'transf_decoder._decoder.layers.6.layer_norm_2.bias', 'transf_decoder._decoder.layers.6.layer_norm_2.weight', 'transf_decoder._decoder.layers.6.layer_norm_3.bias', 'transf_decoder._decoder.layers.6.layer_norm_3.weight', 'transf_decoder._decoder.layers.6.second_sub_layer.key_net.bias', 'transf_decoder._decoder.layers.6.second_sub_layer.key_net.weight', 'transf_decoder._decoder.layers.6.second_sub_layer.out_projection.bias', 'transf_decoder._decoder.layers.6.second_sub_layer.out_projection.weight', 'transf_decoder._decoder.layers.6.second_sub_layer.query_net.bias', 'transf_decoder._decoder.layers.6.second_sub_layer.query_net.weight', 'transf_decoder._decoder.layers.6.second_sub_layer.value_net.bias', 'transf_decoder._decoder.layers.6.second_sub_layer.value_net.weight', 'transf_decoder._decoder.layers.6.third_sub_layer.dense_in.bias', 'transf_decoder._deco


Error: File content (12515 tokens) exceeds maximum allowed tokens (10000). Use offset and limit parameters to read specific portions of the file, or search for specific content instead of reading the whole file.


11:Model loaded in 6.4s
18:Traceback (most recent call last):
39:RuntimeError: Expected a 'cuda' device type for generator but found 'cpu'


14	Filtered to 1000 samples for languages: ['hi', 'en']
15	Dataset: 1000 samples
16
17	Running inference (batch_size=1)...
18	Traceback (most recent call last):
19	  File "/home/ubuntu/training/benchmark_cohere_transcribe.py", line 674, in <module>
20	    main()
21	  File "/home/ubuntu/training/benchmark_cohere_transcribe.py", line 439, in main
22	    results = run_inference(
23	  File "/home/ubuntu/training/benchmark_cohere_transcribe.py", line 327, in run_inference
24	    features = processor.feature_extractor(
25	  File "/home/ubuntu/.cache/huggingface/modules/transformers_modules/cohere_hyphen_transcribe_hyphen_ckpt_hyphen_10000/processing_cohere_asr.py", line 436, in __call__
26	    input_features, length = self.filterbank(audio_tensor, seq_len)
27	  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
28	    return self._call_impl(*args, **kwargs)
29	  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
30	    return forward_call(*args, **kwargs)
31	  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
32	    return func(*args, **kwargs)
33	  File "/home/ubuntu/.cache/huggingface/modules/transformers_modules/cohere_hyphen_transcribe_hyphen_ckpt_hyphen_10000/processing_cohere_asr.py", line 256, in forward
34	    x = self._apply_dither(x, seq_len_time)
35	  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 838, in _fn
36	    return fn(*args, **kwargs)
37	  File "/home/ubuntu/.cache/huggingface/modules/transformers_modules/cohere_hyphen_transcribe_hyphen_ckpt_hyphen_10000/processing_cohere_asr.py", line 166, in _apply_dither
38	    noise = torch.randn(
39	RuntimeError: Expected a 'cuda' device type for generator but found 'cpu'
40

Filtered to 1000 samples for languages: ['hi', 'en']
Dataset: 1000 samples

Running inference (batch_size=1)...
Traceback (most recent call last):
  File "/home/ubuntu/training/benchmark_cohere_transcribe.py", line 674, in <module>
    main()
  File "/home/ubuntu/training/benchmark_cohere_transcribe.py", line 439, 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 dither function expects CUDA tensors but is getting CPU. The issue is that the filterbank is on CPU. Let me fix the feature extractor device setup.


Edit

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

    # Load feature extractor from checkpoint (has custom processing code)
    feature_extractor = AutoFeatureExtractor.from_pretrained(
        args.checkpoint, trust_remote_code=True, 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)
    feature_extractor = AutoFeatureExtractor.from_pretrained(
        args.checkpoint, trust_remote_code=True, 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
    )
    # Move filterbank to GPU for faster feature extraction
    feature_extractor._device = args.device

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    error_analysis = compute_error_analysis(samples_for_errors)

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

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

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

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

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

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

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

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

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

if __name__ == "__main__":
    main()

-        args.checkpoint, trust_remote_code=True

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

-    # Move filterbank to GPU for faster feature extraction

-    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 on Hindi only


if you definitely need something for sure, i can get them either updated to R2 if big files, else git if code files.


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]


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


The R2 checkpoint at `ckpt-10000/` has `model/` and `processor/` subdirectories but is missing a few files needed for inference:

**Missing from R2 (code files - git push would work):**
1. `modeling_cohere_asr.py` — the model's `config.json` auto_map references it for `AutoModelForSpeechSeq2Seq`
2. `tokenization_cohere_asr.py` — referenced in auto_map for `AutoTokenizer`
3. `special_tokens_map.json` — needed by `load_extended_tokenizer()` in `tokenizer_utils.py`

I copied these from `tokenizer_extension/extended_model/` as a stopgap, but those may not match this checkpoint exactly. Could you get the correct versions pushed to R2 under `ckpt-10000/model/` (or `processor/`), or confirm the extended_model versions are correct to use?

Also — there was a runtime error: the feature extractor's dither function expects CUDA tensors. The `AutoFeatureExtractor.from_pretrained()` doesn't seem to accept a `device=` kwarg cleanly. Was there a specific inference script or pattern you used for this model that handles the GPU placement?


Actually, thinking about it more — the files in `tokenizer_extension/extended_model/` are almost certainly correct since that's the code the training script uses, and only the weights change during training. The copies I made should be fine.

For the CUDA device error, let me look at the dither implementation to fix it properly.


147-        self.fb = self.fb.to(dtype=torch.bfloat16)
148-        self.generator = torch.Generator(device=device)
149-        self.generator.manual_seed(0)
150-
151-    @_dynamo_disable
152:    def _apply_dither(self, x, seq_len_time):
153-        """Apply deterministic per-sample dither outside torch.compile.
154-
155-        Each sample is seeded by its valid waveform length so that dither noise
156-        is batch-composition invariant (a sample's features depend only on its
157-        own content, not on what else is in the batch).
--
251-        if self.stft_pad_amount is not None:
252-            x = torch.nn.functional.pad(
253-                x.unsqueeze(1), (self.stft_pad_amount, self.stft_pad_amount), "constant"
254-            ).squeeze(1)
255-
256:        x = self._apply_dither(x, seq_len_time)
257-
258-        if self.preemph is not None:
259-            timemask = torch.arange(x.shape[1], device=x.device).unsqueeze(0) < seq_len_time.unsqueeze(1)
260-            x = torch.cat((x[:, 0].unsqueeze(1), x[:, 1:] - self.preemph * x[:, :-1]), dim=1)
261-            x = x.masked_fill(~timemask, 0.0)


The issue is clear: `self.generator = torch.Generator(device=device)` — when `from_pretrained` creates the feature extractor, `device` defaults to `"cpu"`, so the generator is on CPU. But the audio tensors are sent to CUDA, causing the mismatch.

The fix is to ensure the filterbank is initialized with the right device. Let me update the benchmark script.