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

            # Construct prompt — use <audio_soft_token> which processor expands to full audio sequence
            user_content = f"<audio_soft_token>\nTranscribe the audio.\nLanguage: {lang_name}"

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

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

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

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

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

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

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

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

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

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

    return wav

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

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

        for i, sample in enumerate(batch):
            lang_name = LANG_CODE_TO_NAME.get(sample["lang_code"], sample["language"].capitalize())

            # Construct prompt matching training format
            user_content = f"<audio>\nTranscribe the audio.\nLanguage: {lang_name}"
            messages = [{"role": "user", "content": user_content}]

            # Get audio array
            wav = load_audio_array(sample)

            try:
                # Apply chat template
                prompt_text = processor.apply_chat_template(
                    messages, add_generation_prompt=True, tokenize=False
                )

                # Process inputs (text + audio)
                inputs = processor(
                    text=prompt_text,
                    audio=[wav],
                    sampling_rate=16000,
                    return_tensors="pt",
                )

                # Move to model device
                inputs = {k: v.to(model.device) if isinstance(v, torch.Tensor) else v
                          for k, v in inputs.items()}

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

                # Decode only the new tokens
                input_len = inputs["input_ids"].shape[1]
                generated_ids = output_ids[0][input_len:]
                hypothesis = processor.tokenizer.decode(generated_ids, skip_special_tokens=True).strip()

            except Exception as e:
                print(f"  Sample {start + i} ({sample['id']}) FAILED: {e}")
                hypothesis = ""

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

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

    return results

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

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

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

        if not ref_norm:
            continue

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

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

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

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

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

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

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

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

    return metrics

def print_metrics_table(metrics: dict, checkpoint_name: str):
    """Pretty print metrics as a table."""
    print(f"\n{'='*85}")
    print(f"  BENCHMARK RESULTS: {checkpoint_name}")
    print(f"{'='*85}")
    print(f"{'Language':<14} {'N':>5} {'WER%':>8} {'CER%':>8} {'WER-n%':>8} {'CER-n%':>8}")
    print(f"{'-'*14} {'-'*5} {'-'*8} {'-'*8} {'-'*8} {'-'*8}")

    for lang in sorted(k for k in metrics if not k.startswith("__")):
        m = metrics[lang]
        print(f"{lang:<14} {m['n_samples']:>5} {m['wer']:>8.2f} {m['cer']:>8.2f} "
              f"{m['wer_normalized']:>8.2f} {m['cer_normalized']:>8.2f}")

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

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

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

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

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

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

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

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

    # Load model and processor
    print(f"Loading checkpoint: {args.checkpoint}")
    t0 = time.time()

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

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

    processor = AutoProcessor.from_pretrained(
        args.checkpoint, trust_remote_code=True, tokenizer=tokenizer
    )
    model = AutoModelForImageTextToText.from_pretrained(
        args.checkpoint,
        torch_dtype=torch.bfloat16,
        device_map=args.device,
        trust_remote_code=True,
    )
    model.eval()

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

    # Load dataset
    hf_token = args.hf_token or os.environ.get("HF_TOKEN")
    print(f"Loading dataset: {args.hf_dataset}")
    dataset = hf_load_dataset(
        args.hf_dataset, token=hf_token, cache_dir=args.cache_dir, split="train"
    )

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

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

    # Run inference
    print(f"\nRunning inference (batch_size={args.batch_size})...")
    t0 = time.time()
    results = run_inference(model, processor, dataset, args.batch_size, args.max_new_tokens)
    inference_time = time.time() - t0

    total_audio_sec = sum(r["duration"] for r in results)
    rtf = inference_time / total_audio_sec if total_audio_sec > 0 else 0
    print(f"Inference done in {inference_time:.1f}s for {total_audio_sec:.0f}s audio (RTF={rtf:.4f})")

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

    # Compute metrics
    metrics = compute_metrics(results)
    metrics["__meta__"] = {
        "checkpoint": args.checkpoint,
        "checkpoint_name": ckpt_name,
        "model_type": "gemma3n-E2B-asr",
        "dataset": args.hf_dataset,
        "batch_size": args.batch_size,
        "inference_time_sec": round(inference_time, 2),
        "total_audio_sec": round(total_audio_sec, 2),
        "rtf": round(rtf, 4),
    }

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

    # Print table
    print_metrics_table(metrics, ckpt_name)

    # Save sample analysis
    save_sample_analysis(results, output_dir)

if __name__ == "__main__":
    main()

-            # Construct prompt matching training format

-            user_content = f"<audio>\nTranscribe the audio.\nLanguage: {lang_name}"

+            # Construct prompt — use <audio_soft_token> which processor expands to full audio sequence

+            user_content = f"<audio_soft_token>\nTranscribe the audio.\nLanguage: {lang_name}"


Now let me kill the failing run and quick-test this fix.


Bash

# Quick single-sample test with the fix
python3 -c "
import numpy as np
import torch
from transformers import AutoTokenizer, AutoConfig, AutoProcessor, AutoModelForImageTextToText
from datasets import load_dataset

ckpt = '/home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-10000'
config = AutoConfig.from_pretrained(ckpt, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(ckpt, trust_remote_code=True)

patches = {
    'audio_token_id': config.audio_token_id, 'audio_token': '<audio_soft_token>',
    'boa_token_id': config.boa_token_id, 'boa_token': '<start_of_audio>',
    'eoa_token_id': config.eoa_token_id, 'eoa_token': '<end_of_audio>',
    'boi_token_id': config.boi_token_id, 'boi_token': '<start_of_image>',
    'eoi_token_id': config.eoi_token_id, 'eoi_token': '<end_of_image>',
    'image_token_id': config.image_token_id, 'image_token': '<image_soft_token>',
}
for k, v in patches.items():
    if not hasattr(tokenizer, k):
        setattr(tokenizer, k, v)

processor = AutoProcessor.from_pretrained(ckpt, trust_remote_code=True, tokenizer=tokenizer)

model = AutoModelForImageTextToText.from_pretrained(
    ckpt, dtype=torch.bfloat16, device_map='cuda:0', trust_remote_code=True
)
model.eval()

# Load first Hindi sample
ds = load_dataset('BayAreaBoys/indic-asr-benchmark-6k',
    cache_dir='/home/ubuntu/training/datasets/indic-asr-benchmark-6k', split='train')
ds = ds.filter(lambda x: x['lang_code'] == 'hi')
sample = ds[0]

wav = np.array(sample['audio']['array'], dtype=np.float32)
sr = sample['audio']['sampling_rate']
if sr != 16000:
    import torchaudio
    wav = torchaudio.functional.resample(torch.from_numpy(wav).unsqueeze(0), sr, 16000).squeeze(0).numpy()

messages = [{'role': 'user', 'content': '<audio_soft_token>\nTranscribe the audio.\nLanguage: Hindi'}]
prompt = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
print('PROMPT (first 200):', repr(prompt[:200]))

inputs = processor(text=prompt, audio=[wav], sampling_rate=16000, return_tensors='pt')
print('input_ids shape:', inputs['input_ids'].shape)

# Check audio tokens
audio_token_count = (inputs['input_ids'] == config.audio_token_id).sum().item()
print(f'Audio token count in input: {audio_token_count}')

inputs = {k: v.to(model.device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}

with torch.no_grad():
    out = model.generate(**inputs, max_new_tokens=256, do_sample=False)

gen_ids = out[0][inputs['input_ids'].shape[1]:]
text = processor.tokenizer.decode(gen_ids, skip_special_tokens=True).strip()
print(f'REF: {sample[\"reference\"]}')
print(f'HYP: {text}')
" 2>&1

Single-sample test with audio token fix


Some weights of the model checkpoint at /home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-10000 were not used when initializing Gemma3nForConditionalGeneration: ['model.vision_tower.timm_model.conv_stem.conv.bias']
- This IS expected if you are initializing Gemma3nForConditionalGeneration 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 Gemma3nForConditionalGeneration from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
Using the latest cached version of the dataset since BayAreaBoys/indic-asr-benchmark-6k couldn't be found on the Hugging Face Hub
WARNING:datasets.load:Using the latest cached version of the dataset since BayAreaBoys/indic-asr-benchmark-6k couldn't be found on the Hugging Face Hub
Found the latest cached dataset configuration 'default' at /home/ubuntu/training/datasets/indic-asr-benchmark-6k/BayAreaBoys___indic-asr-benchmark-6k/default/0.0.0/4d4bc89b7c915a6d80a1efbee3a708d58e688c81 (last modified on Thu Mar 26 00:59:49 2026).
WARNING:datasets.packaged_modules.cache.cache:Found the latest cached dataset configuration 'default' at /home/ubuntu/training/datasets/indic-asr-benchmark-6k/BayAreaBoys___indic-asr-benchmark-6k/default/0.0.0/4d4bc89b7c915a6d80a1efbee3a708d58e688c81 (last modified on Thu Mar 26 00:59:49 2026).

Filter:   0%|          | 0/6000 [00:00<?, ? examples/s]
Filter:  17%|█▋        | 1000/6000 [00:06<00:31, 158.94 examples/s]
Filter:  33%|███▎      | 2000/6000 [00:12<00:24, 163.05 examples/s]
Filter:  50%|█████     | 3000/6000 [00:19<00:19, 154.86 examples/s]
Filter:  67%|██████▋   | 4000/6000 [00:26<00:13, 150.37 examples/s]
Filter:  83%|████████▎ | 5000/6000 [00:32<00:06, 149.76 examples/s]
Filter: 100%|██████████| 6000/6000 [00:39<00:00, 149.29 examples/s]
Filter: 100%|██████████| 6000/6000 [00:39<00:00, 151.62 examples/s]
The following generation flags are not valid and may be ignored: ['top_p', 'top_k']. Set `TRANSFORMERS_VERBOSITY=info` for more details.
PROMPT (first 200): '<bos><start_of_turn>user\n<audio_soft_token>\nTranscribe the audio.\nLanguage: Hindi<end_of_turn>\n<start_of_turn>model\n'
input_ids shape: torch.Size([1, 210])
Audio token count in input: 188
REF: इस तरह इसका सोलहदिवसीय क्रम लगातार जारी है
HYP: इस तरह इसका सोलह दिवसीय क्रम लगातार जारी है


Excellent! It's working now! The transcription is almost identical to the reference:
- **REF:** `इस तरह इसका सोलहदिवसीय क्रम लगातार जारी है`
- **HYP:** `इस तरह इसका सोलह दिवसीय क्रम लगातार जारी है`

Only difference is a space in `सोलहदिवसीय` vs `सोलह दिवसीय`. The model is producing meaningful transcriptions!

Now let me run the full benchmark on all 6000 samples. Let me first kill the failing background run.


pkill -f "benchmark_gemma3n_asr.*ckpt-10000-smoke" 2>&1; echo "Killed old run"

Kill the failing smoke test


Exit code 144


python3 benchmark_gemma3n_asr.py \
    --checkpoint /home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-10000 \
    --checkpoint-name ckpt-10000 \
    --output-dir /home/ubuntu/training/benchmark_results_gemma3n \
    --batch-size 1 2>&1

Full benchmark on ckpt-10000 (6000 samples)


Command running in background with ID: b9ofklw8a. Output is being written to: /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/b9ofklw8a.output


Let me check the ckpt-20000 download progress and monitor the benchmark.


du -sh /home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-20000/ 2>&1 && ls -lh /home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-20000/ 2>&1

Check ckpt-20000 download progress


21G	/home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-20000/
total 21G
-rw-rw-r-- 1 ubuntu ubuntu  613 Mar 25 22:59 chat_template.jinja
-rw-rw-r-- 1 ubuntu ubuntu 4.5K Mar 25 22:59 config.json
-rw-rw-r-- 1 ubuntu ubuntu  209 Mar 25 22:59 generation_config.json
-rw-rw-r-- 1 ubuntu ubuntu  21G Mar 25 23:29 model.safetensors
-rw-rw-r-- 1 ubuntu ubuntu 2.3K Mar 25 22:59 processor_config.json
-rw-rw-r-- 1 ubuntu ubuntu 1004 Mar 25 22:59 tokenizer_config.json
-rw-rw-r-- 1 ubuntu ubuntu  32M Mar 25 22:59 tokenizer.json


ckpt-20000 download is complete (21GB model only, skipped optimizer). Let me check benchmark progress.


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/b9ofklw8a.output 2>&1

Check ckpt-10000 benchmark progress


Using the latest cached version of the dataset since BayAreaBoys/indic-asr-benchmark-6k couldn't be found on the Hugging Face Hub
WARNING:datasets.load:Using the latest cached version of the dataset since BayAreaBoys/indic-asr-benchmark-6k couldn't be found on the Hugging Face Hub
Found the latest cached dataset configuration 'default' at /home/ubuntu/training/datasets/indic-asr-benchmark-6k/BayAreaBoys___indic-asr-benchmark-6k/default/0.0.0/4d4bc89b7c915a6d80a1efbee3a708d58e688c81 (last modified on Thu Mar 26 01:11:34 2026).
WARNING:datasets.packaged_modules.cache.cache:Found the latest cached dataset configuration 'default' at /home/ubuntu/training/datasets/indic-asr-benchmark-6k/BayAreaBoys___indic-asr-benchmark-6k/default/0.0.0/4d4bc89b7c915a6d80a1efbee3a708d58e688c81 (last modified on Thu Mar 26 01:11:34 2026).
The following generation flags are not valid and may be ignored: ['top_p', 'top_k']. Set `TRANSFORMERS_VERBOSITY=info` for more details.


Still loading. Let me wait a bit more and check again.


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/b9ofklw8a.output 2>&1

Check benchmark progress


Running inference (batch_size=1)...
  [10/6000] (0%) processed


Good, it's running. At the smoke test rate (~500 samples in 262s), 6000 samples should take about ~30 minutes. Let me check progress periodically. Meanwhile, let me also look at the print frequency — currently it prints every 10 samples but only when `end % max(batch_size * 10, 10) == 0`. That's fine.


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/b9ofklw8a.output 2>&1

Check benchmark progress


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/b9ofklw8a.output 2>&1

Check benchmark progress


Running inference (batch_size=1)...
  [10/6000] (0%) processed
  [20/6000] (0%) processed


The benchmark is running. At ~2 samples/sec, the full 6000-sample run should take about 50 minutes. Let me check back periodically.


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/b9ofklw8a.output 2>&1

Check benchmark progress


Dataset: 6000 samples

Running inference (batch_size=1)...
  [10/6000] (0%) processed
  [20/6000] (0%) processed


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/b9ofklw8a.output 2>&1

Check progress


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/b9ofklw8a.output 2>&1

Check progress


nvidia-smi --query-compute-apps=pid,name,used_gpu_memory --format=csv 2>&1

Check GPU processes


pid, process_name, used_gpu_memory [MiB]
113020, python3, 18722 MiB


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/b9ofklw8a.output 2>&1

Check progress


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/b9ofklw8a.output 2>&1

Check progress


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/b9ofklw8a.output 2>&1

Check progress


ps aux | grep benchmark_gemma3n | grep -v grep 2>&1

Check if benchmark process is still running


ubuntu    112998  0.0  0.0   7504  3456 ?        Ss   01:13   0:00 /bin/bash -c source /home/ubuntu/.claude/shell-snapshots/snapshot-bash-1774479679611-o2lfid.sh && shopt -u extglob 2>/dev/null || true && eval 'python3 benchmark_gemma3n_asr.py \     --checkpoint /home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-10000 \     --checkpoint-name ckpt-10000 \     --output-dir /home/ubuntu/training/benchmark_results_gemma3n \     --batch-size 1 2>&1' \< /dev/null && pwd -P >| /tmp/claude-8077-cwd
ubuntu    113020  147  1.8 36891188 2299332 ?    RLl  01:13   2:40 python3 benchmark_gemma3n_asr.py --checkpoint /home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-10000 --checkpoint-name ckpt-10000 --output-dir /home/ubuntu/training/benchmark_results_gemma3n --batch-size 1


wc -l /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/b9ofklw8a.output 2>&1

Count output lines


18 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/b9ofklw8a.output


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/b9ofklw8a.output 2>&1

Check latest output


  [10/6000] (0%) processed
  [20/6000] (0%) processed
  [30/6000] (0%) processed


Good, it's progressing — just slowly. The output buffering makes it appear stuck. 30/6000 after ~2 minutes. That's about 15 samples/min. At this rate it'd take ~6.5 hours. That's too slow.

Let me check if there's a way to speed this up. The model is 18GB in memory on a 80GB GPU — plenty of room. The issue is probably that we're processing one sample at a time and the model has to process 188 audio tokens per sample plus the text. Let me check if I can batch the inputs.

Actually, the bigger issue might be the dataset loading — each sample requires decoding audio which is slow. Let me check the actual throughput.


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/b9ofklw8a.output 2>&1

Check latest progress


Running inference (batch_size=1)...
  [10/6000] (0%) processed
  [20/6000] (0%) processed
  [30/6000] (0%) processed


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/b9ofklw8a.output 2>&1

Check progress


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/b9ofklw8a.output 2>&1

Check progress


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/b9ofklw8a.output 2>&1

Check progress


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/b9ofklw8a.output 2>&1

Check progress


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/b9ofklw8a.output 2>&1

Check progress


Running inference (batch_size=1)...
  [10/6000] (0%) processed
  [20/6000] (0%) processed
  [30/6000] (0%) processed
  [40/6000] (1%) processed


At 10 samples/min, 6000 would take ~10 hours. Too slow. The model config has `use_cache: false`. Let me kill this and optimize.


kill 113020 2>&1; echo "Killed benchmark process"

Kill slow benchmark


Killed benchmark process


<task-notification>
<task-id>b9ofklw8a</task-id>
<tool-use-id>toolu_0131Ac4JkqoU186adoBiNHzG</tool-use-id>
<output-file>/tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/b9ofklw8a.output</output-file>
<status>failed</status>
<summary>Background command "Full benchmark on ckpt-10000 (6000 samples)" failed with exit code 143</summary>
</task-notification>
Read the output file to retrieve the result: /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/b9ofklw8a.output


python3 -c "
from transformers import AutoConfig
config = AutoConfig.from_pretrained('/home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-10000', trust_remote_code=True)
print('use_cache:', config.use_cache)
print('text use_cache:', config.text_config.use_cache if hasattr(config.text_config, 'use_cache') else 'N/A')
" 2>&1

Check use_cache setting


use_cache: False
text use_cache: False


`use_cache` is False — that means no KV cache during generation, which makes autoregressive generation extremely slow since it recomputes all attention at every step. Let me enable it.


   315→    )
   316→    model = AutoModelForImageTextToText.from_pretrained(
   317→        args.checkpoint,
   318→        torch_dtype=torch.bfloat16,
   319→        device_map=args.device,
   320→        trust_remote_code=True,
   321→    )
   322→    model.eval()
   323→
   324→    load_time = time.time() - t0
   325→    print(f"Model loaded in {load_time:.1f}s")
   326→
   327→    # Load dataset
   328→    hf_token = args.hf_token or os.environ.get("HF_TOKEN")
   329→    print(f"Loading dataset: {args.hf_dataset}")

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

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

    # Load dataset
    hf_token = args.hf_token or os.environ.get("HF_TOKEN")
    print(f"Loading dataset: {args.hf_dataset}")


Edit

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

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


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

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

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

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

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

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

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

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

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

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

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

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

    return wav

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

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

        for i, sample in enumerate(batch):
            lang_name = LANG_CODE_TO_NAME.get(sample["lang_code"], sample["language"].capitalize())

            # Construct prompt — use <audio_soft_token> which processor expands to full audio sequence
            user_content = f"<audio_soft_token>\nTranscribe the audio.\nLanguage: {lang_name}"
            messages = [{"role": "user", "content": user_content}]

            # Get audio array
            wav = load_audio_array(sample)

            try:
                # Apply chat template
                prompt_text = processor.apply_chat_template(
                    messages, add_generation_prompt=True, tokenize=False
                )

                # Process inputs (text + audio)
                inputs = processor(
                    text=prompt_text,
                    audio=[wav],
                    sampling_rate=16000,
                    return_tensors="pt",
                )

                # Move to model device
                inputs = {k: v.to(model.device) if isinstance(v, torch.Tensor) else v
                          for k, v in inputs.items()}

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

                # Decode only the new tokens
                input_len = inputs["input_ids"].shape[1]
                generated_ids = output_ids[0][input_len:]
                hypothesis = processor.tokenizer.decode(generated_ids, skip_special_tokens=True).strip()

            except Exception as e:
                print(f"  Sample {start + i} ({sample['id']}) FAILED: {e}")
                hypothesis = ""

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

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

    return results

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

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

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

        if not ref_norm:
            continue

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

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

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

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

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

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

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

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

    return metrics

def print_metrics_table(metrics: dict, checkpoint_name: str):
    """Pretty print metrics as a table."""
    print(f"\n{'='*85}")
    print(f"  BENCHMARK RESULTS: {checkpoint_name}")
    print(f"{'='*85}")
    print(f"{'Language':<14} {'N':>5} {'WER%':>8} {'CER%':>8} {'WER-n%':>8} {'CER-n%':>8}")
    print(f"{'-'*14} {'-'*5} {'-'*8} {'-'*8} {'-'*8} {'-'*8}")

    for lang in sorted(k for k in metrics if not k.startswith("__")):
        m = metrics[lang]
        print(f"{lang:<14} {m['n_samples']:>5} {m['wer']:>8.2f} {m['cer']:>8.2f} "
              f"{m['wer_normalized']:>8.2f} {m['cer_normalized']:>8.2f}")

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

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

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

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

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

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

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

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

    # Load model and processor
    print(f"Loading checkpoint: {args.checkpoint}")
    t0 = time.time()

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

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

    processor = AutoProcessor.from_pretrained(
        args.checkpoint, trust_remote_code=True, tokenizer=tokenizer
    )
    model = AutoModelForImageTextToText.from_pretrained(
        args.checkpoint,
        torch_dtype=torch.bfloat16,
        device_map=args.device,
        trust_remote_code=True,
    )
    model.eval()

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

    # Load dataset
    hf_token = args.hf_token or os.environ.get("HF_TOKEN")
    print(f"Loading dataset: {args.hf_dataset}")
    dataset = hf_load_dataset(
        args.hf_dataset, token=hf_token, cache_dir=args.cache_dir, split="train"
    )

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

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

    # Run inference
    print(f"\nRunning inference (batch_size={args.batch_size})...")
    t0 = time.time()
    results = run_inference(model, processor, dataset, args.batch_size, args.max_new_tokens)
    inference_time = time.time() - t0

    total_audio_sec = sum(r["duration"] for r in results)
    rtf = inference_time / total_audio_sec if total_audio_sec > 0 else 0
    print(f"Inference done in {inference_time:.1f}s for {total_audio_sec:.0f}s audio (RTF={rtf:.4f})")

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

    # Compute metrics
    metrics = compute_metrics(results)
    metrics["__meta__"] = {
        "checkpoint": args.checkpoint,
        "checkpoint_name": ckpt_name,
        "model_type": "gemma3n-E2B-asr",
        "dataset": args.hf_dataset,
        "batch_size": args.batch_size,
        "inference_time_sec": round(inference_time, 2),
        "total_audio_sec": round(total_audio_sec, 2),
        "rtf": round(rtf, 4),
    }

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

    # Print table
    print_metrics_table(metrics, ckpt_name)

    # Save sample analysis
    save_sample_analysis(results, output_dir)

if __name__ == "__main__":
    main()

+    # Enable KV cache for faster autoregressive generation (training disables it)

+    if hasattr(model.config, "text_config"):

+        model.config.text_config.use_cache = True


Now let me also check if the generation_config.json uses `cache_implementation: "hybrid"` — this might need specific support.


Bash

# Quick test: verify use_cache=True works and speeds things up
python3 -c "
import numpy as np, torch, time
from transformers import AutoTokenizer, AutoConfig, AutoProcessor, AutoModelForImageTextToText
from datasets import load_dataset

ckpt = '/home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-10000'
config = AutoConfig.from_pretrained(ckpt, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(ckpt, trust_remote_code=True)
patches = {
    'audio_token_id': config.audio_token_id, 'audio_token': '<audio_soft_token>',
    'boa_token_id': config.boa_token_id, 'boa_token': '<start_of_audio>',
    'eoa_token_id': config.eoa_token_id, 'eoa_token': '<end_of_audio>',
    'boi_token_id': config.boi_token_id, 'boi_token': '<start_of_image>',
    'eoi_token_id': config.eoi_token_id, 'eoi_token': '<end_of_image>',
    'image_token_id': config.image_token_id, 'image_token': '<image_soft_token>',
}
for k, v in patches.items():
    if not hasattr(tokenizer, k):
        setattr(tokenizer, k, v)

processor = AutoProcessor.from_pretrained(ckpt, trust_remote_code=True, tokenizer=tokenizer)
model = AutoModelForImageTextToText.from_pretrained(ckpt, dtype=torch.bfloat16, device_map='cuda:0', trust_remote_code=True)
model.eval()
model.config.use_cache = True
model.config.text_config.use_cache = True

ds = load_dataset('BayAreaBoys/indic-asr-benchmark-6k', cache_dir='/home/ubuntu/training/datasets/indic-asr-benchmark-6k', split='train')
ds = ds.filter(lambda x: x['lang_code'] == 'hi')

# Time 5 samples
times = []
for i in range(5):
    sample = ds[i]
    wav = np.array(sample['audio']['array'], dtype=np.float32)
    sr = sample['audio']['sampling_rate']
    if sr != 16000:
        import torchaudio
        wav = torchaudio.functional.resample(torch.from_numpy(wav).unsqueeze(0), sr, 16000).squeeze(0).numpy()

    messages = [{'role': 'user', 'content': '<audio_soft_token>\nTranscribe the audio.\nLanguage: Hindi'}]
    prompt = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
    inputs = processor(text=prompt, audio=[wav], sampling_rate=16000, return_tensors='pt')
    inputs = {k: v.to(model.device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}

    t0 = time.time()
    with torch.no_grad():
        out = model.generate(**inputs, max_new_tokens=256, do_sample=False)
    times.append(time.time() - t0)

    gen_ids = out[0][inputs['input_ids'].shape[1]:]
    text = processor.tokenizer.decode(gen_ids, skip_special_tokens=True).strip()
    print(f'Sample {i}: {times[-1]:.2f}s | HYP: {text[:60]}...')

print(f'\nAvg: {np.mean(times):.2f}s/sample, Rate: {60/np.mean(times):.0f} samples/min')
" 2>&1

Test with use_cache=True for speedup


Some weights of the model checkpoint at /home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-10000 were not used when initializing Gemma3nForConditionalGeneration: ['model.vision_tower.timm_model.conv_stem.conv.bias']
- This IS expected if you are initializing Gemma3nForConditionalGeneration 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 Gemma3nForConditionalGeneration from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
Using the latest cached version of the dataset since BayAreaBoys/indic-asr-benchmark-6k couldn't be found on the Hugging Face Hub
WARNING:datasets.load:Using the latest cached version of the dataset since BayAreaBoys/indic-asr-benchmark-6k couldn't be found on the Hugging Face Hub
Found the latest cached dataset configuration 'default' at /home/ubuntu/training/datasets/indic-asr-benchmark-6k/BayAreaBoys___indic-asr-benchmark-6k/default/0.0.0/4d4bc89b7c915a6d80a1efbee3a708d58e688c81 (last modified on Thu Mar 26 01:11:34 2026).
WARNING:datasets.packaged_modules.cache.cache:Found the latest cached dataset configuration 'default' at /home/ubuntu/training/datasets/indic-asr-benchmark-6k/BayAreaBoys___indic-asr-benchmark-6k/default/0.0.0/4d4bc89b7c915a6d80a1efbee3a708d58e688c81 (last modified on Thu Mar 26 01:11:34 2026).
The following generation flags are not valid and may be ignored: ['top_p', 'top_k']. Set `TRANSFORMERS_VERBOSITY=info` for more details.
Sample 0: 1.52s | HYP: इस तरह इसका सोलह दिवसीय क्रम लगातार जारी है...
Sample 1: 1.89s | HYP: इसके अलावा ऑल राउंडर कोरी एंडरसन और ग्रांट एलियट भी लय में ह...
Sample 2: 0.84s | HYP: हमारी सरकार ने सिंचाई की व्यवस्था की है...
Sample 3: 1.54s | HYP: पुणे और मुंबई का अंतरराष्ट्रीय हवाई अड्डा अनुक्रमे एक सौ बीस...
Sample 4: 1.37s | HYP: पिछले कुछ दिनों से व्यापारियों के पास वसूली के लिए कॉल आ रहा...

Avg: 1.43s/sample, Rate: 42 samples/min


Huge speedup! From ~10 samples/min to **42 samples/min** with KV cache enabled. The transcriptions look great too — very clean Hindi text.

At 42 samples/min, 6000 samples = ~2.4 hours. Still slow. Let me also check if we can get flash attention working.


python3 -c "
from transformers import AutoModelForImageTextToText
import torch
# Test with flash attention
model = AutoModelForImageTextToText.from_pretrained(
    '/home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-10000',
    dtype=torch.bfloat16, device_map='cuda:0', trust_remote_code=True,
    attn_implementation='flash_attention_2'
)
print('Loaded with flash attention 2')
" 2>&1

Test flash attention support


Error: Exit code 1
Traceback (most recent call last):
  File "<string>", line 5, in <module>
  File "/home/ubuntu/.local/lib/python3.10/site-packages/transformers/models/auto/auto_factory.py", line 604, in from_pretrained
    return model_class.from_pretrained(
  File "/home/ubuntu/.local/lib/python3.10/site-packages/transformers/modeling_utils.py", line 277, in _wrapper
    return func(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/transformers/modeling_utils.py", line 4971, in from_pretrained
    model = cls(config, *model_args, **model_kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/transformers/models/gemma3n/modeling_gemma3n.py", line 2175, in __init__
    super().__init__(config)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/transformers/modeling_utils.py", line 2076, in __init__
    self.config._attn_implementation_internal = self._check_and_adjust_attn_implementation(
  File "/home/ubuntu/.local/lib/python3.10/site-packages/transformers/modeling_utils.py", line 2686, in _check_and_adjust_attn_implementation
    applicable_attn_implementation = self.get_correct_attn_implementation(
  File "/home/ubuntu/.local/lib/python3.10/site-packages/transformers/modeling_utils.py", line 2714, in get_correct_attn_implementation
    self._flash_attn_2_can_dispatch(is_init_check)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/transformers/modeling_utils.py", line 2422, in _flash_attn_2_can_dispatch
    raise ImportError(f"{preface} the package flash_attn seems to be not installed. {install_message}")
ImportError: FlashAttention2 has been toggled on, but it cannot be used due to the following error: the package flash_attn seems to be not installed. Please refer to the documentation of https://huggingface.co/docs/transformers/perf_infer_gpu_one#flashattention-2 to install Flash Attention 2.


No flash attention available. Let me try `sdpa` instead.


python3 -c "
from transformers import AutoModelForImageTextToText
import torch
model = AutoModelForImageTextToText.from_pretrained(
    '/home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-10000',
    dtype=torch.bfloat16, device_map='cuda:0', trust_remote_code=True,
    attn_implementation='sdpa'
)
print('Loaded with SDPA attention')
del model
" 2>&1

Test SDPA attention support


Some weights of the model checkpoint at /home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-10000 were not used when initializing Gemma3nForConditionalGeneration: ['model.vision_tower.timm_model.conv_stem.conv.bias']
- This IS expected if you are initializing Gemma3nForConditionalGeneration 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 Gemma3nForConditionalGeneration from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
Loaded with SDPA attention


Good. Let me add SDPA support and also add `attn_implementation` as a CLI arg. At 42 samples/min = ~2.4 hours. That's workable but let me also add `--attn-implementation` option and see if SDPA helps.

Actually, 2.4 hours is fine. Let me just kick off the full run now with the KV cache fix, and also add SDPA.


Edit

    model = AutoModelForImageTextToText.from_pretrained(
        args.checkpoint,
        torch_dtype=torch.bfloat16,
        device_map=args.device,
        trust_remote_code=True,
    )

    model = AutoModelForImageTextToText.from_pretrained(
        args.checkpoint,
        torch_dtype=torch.bfloat16,
        device_map=args.device,
        trust_remote_code=True,
        attn_implementation="sdpa",
    )


    model = AutoModelForImageTextToText.from_pretrained(
        args.checkpoint,
        torch_dtype=torch.bfloat16,
        device_map=args.device,
        trust_remote_code=True,
        attn_implementation="sdpa",
    )

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

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

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

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

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

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

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

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

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

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

    return wav

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

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

        for i, sample in enumerate(batch):
            lang_name = LANG_CODE_TO_NAME.get(sample["lang_code"], sample["language"].capitalize())

            # Construct prompt — use <audio_soft_token> which processor expands to full audio sequence
            user_content = f"<audio_soft_token>\nTranscribe the audio.\nLanguage: {lang_name}"
            messages = [{"role": "user", "content": user_content}]

            # Get audio array
            wav = load_audio_array(sample)

            try:
                # Apply chat template
                prompt_text = processor.apply_chat_template(
                    messages, add_generation_prompt=True, tokenize=False
                )

                # Process inputs (text + audio)
                inputs = processor(
                    text=prompt_text,
                    audio=[wav],
                    sampling_rate=16000,
                    return_tensors="pt",
                )

                # Move to model device
                inputs = {k: v.to(model.device) if isinstance(v, torch.Tensor) else v
                          for k, v in inputs.items()}

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

                # Decode only the new tokens
                input_len = inputs["input_ids"].shape[1]
                generated_ids = output_ids[0][input_len:]
                hypothesis = processor.tokenizer.decode(generated_ids, skip_special_tokens=True).strip()

            except Exception as e:
                print(f"  Sample {start + i} ({sample['id']}) FAILED: {e}")
                hypothesis = ""

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

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

    return results

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

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

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

        if not ref_norm:
            continue

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

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

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

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

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

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

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

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

    return metrics

def print_metrics_table(metrics: dict, checkpoint_name: str):
    """Pretty print metrics as a table."""
    print(f"\n{'='*85}")
    print(f"  BENCHMARK RESULTS: {checkpoint_name}")
    print(f"{'='*85}")
    print(f"{'Language':<14} {'N':>5} {'WER%':>8} {'CER%':>8} {'WER-n%':>8} {'CER-n%':>8}")
    print(f"{'-'*14} {'-'*5} {'-'*8} {'-'*8} {'-'*8} {'-'*8}")

    for lang in sorted(k for k in metrics if not k.startswith("__")):
        m = metrics[lang]
        print(f"{lang:<14} {m['n_samples']:>5} {m['wer']:>8.2f} {m['cer']:>8.2f} "
              f"{m['wer_normalized']:>8.2f} {m['cer_normalized']:>8.2f}")

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

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

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

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

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

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

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

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

    # Load model and processor
    print(f"Loading checkpoint: {args.checkpoint}")
    t0 = time.time()

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

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

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

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

    # Load dataset
    hf_token = args.hf_token or os.environ.get("HF_TOKEN")
    print(f"Loading dataset: {args.hf_dataset}")
    dataset = hf_load_dataset(
        args.hf_dataset, token=hf_token, cache_dir=args.cache_dir, split="train"
    )

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

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

    # Run inference
    print(f"\nRunning inference (batch_size={args.batch_size})...")
    t0 = time.time()
    results = run_inference(model, processor, dataset, args.batch_size, args.max_new_tokens)
    inference_time = time.time() - t0

    total_audio_sec = sum(r["duration"] for r in results)
    rtf = inference_time / total_audio_sec if total_audio_sec > 0 else 0
    print(f"Inference done in {inference_time:.1f}s for {total_audio_sec:.0f}s audio (RTF={rtf:.4f})")

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

    # Compute metrics
    metrics = compute_metrics(results)
    metrics["__meta__"] = {
        "checkpoint": args.checkpoint,
        "checkpoint_name": ckpt_name,
        "model_type": "gemma3n-E2B-asr",
        "dataset": args.hf_dataset,
        "batch_size": args.batch_size,
        "inference_time_sec": round(inference_time, 2),
        "total_audio_sec": round(total_audio_sec, 2),
        "rtf": round(rtf, 4),
    }

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

    # Print table
    print_metrics_table(metrics, ckpt_name)

    # Save sample analysis
    save_sample_analysis(results, output_dir)

if __name__ == "__main__":
    main()


Bash

python3 benchmark_gemma3n_asr.py \
    --checkpoint /home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-10000 \
    --checkpoint-name ckpt-10000 \
    --output-dir /home/ubuntu/training/benchmark_results_gemma3n \
    --batch-size 1 2>&1

Full benchmark ckpt-10000 with KV cache + SDPA


Command running in background with ID: brrrsmz06. Output is being written to: /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check new benchmark progress


(Bash completed with no output)


tail -10 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check new benchmark progress


Read full output


ls -la /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check output file


-rw-rw-r-- 1 ubuntu ubuntu 0 Mar 26 01:18 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output


ps aux | grep benchmark_gemma3n | grep -v grep 2>&1

Check running processes


ubuntu    115714  0.0  0.0   7504  3456 ?        Ss   01:18   0:00 /bin/bash -c source /home/ubuntu/.claude/shell-snapshots/snapshot-bash-1774479679611-o2lfid.sh && shopt -u extglob 2>/dev/null || true && eval 'python3 benchmark_gemma3n_asr.py \     --checkpoint /home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-10000 \     --checkpoint-name ckpt-10000 \     --output-dir /home/ubuntu/training/benchmark_results_gemma3n \     --batch-size 1 2>&1' \< /dev/null && pwd -P >| /tmp/claude-2a52-cwd
ubuntu    115734  120 12.8 53952396 15921524 ?   Rl   01:18   0:25 python3 benchmark_gemma3n_asr.py --checkpoint /home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-10000 --checkpoint-name ckpt-10000 --output-dir /home/ubuntu/training/benchmark_results_gemma3n --batch-size 1


Process is running, just loading the model (using 15.9GB RAM already). Let me wait for it to start producing output.


Check output


`torch_dtype` is deprecated! Use `dtype` instead!
Some weights of the model checkpoint at /home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-10000 were not used when initializing Gemma3nForConditionalGeneration: ['model.vision_tower.timm_model.conv_stem.conv.bias']
- This IS expected if you are initializing Gemma3nForConditionalGeneration 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 Gemma3nForConditionalGeneration from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
Using the latest cached version of the dataset since BayAreaBoys/indic-asr-benchmark-6k couldn't be found on the Hugging Face Hub
WARNING:datasets.load:Using the latest cached version of the dataset since BayAreaBoys/indic-asr-benchmark-6k couldn't be found on the Hugging Face Hub
Found the latest cached dataset configuration 'default' at /home/ubuntu/training/datasets/indic-asr-benchmark-6k/BayAreaBoys___indic-asr-benchmark-6k/default/0.0.0/4d4bc89b7c915a6d80a1efbee3a708d58e688c81 (last modified on Thu Mar 26 01:11:34 2026).
WARNING:datasets.packaged_modules.cache.cache:Found the latest cached dataset configuration 'default' at /home/ubuntu/training/datasets/indic-asr-benchmark-6k/BayAreaBoys___indic-asr-benchmark-6k/default/0.0.0/4d4bc89b7c915a6d80a1efbee3a708d58e688c81 (last modified on Thu Mar 26 01:11:34 2026).
The following generation flags are not valid and may be ignored: ['top_p', 'top_k']. Set `TRANSFORMERS_VERBOSITY=info` for more details.


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check output


Using the latest cached version of the dataset since BayAreaBoys/indic-asr-benchmark-6k couldn't be found on the Hugging Face Hub
WARNING:datasets.load:Using the latest cached version of the dataset since BayAreaBoys/indic-asr-benchmark-6k couldn't be found on the Hugging Face Hub
Found the latest cached dataset configuration 'default' at /home/ubuntu/training/datasets/indic-asr-benchmark-6k/BayAreaBoys___indic-asr-benchmark-6k/default/0.0.0/4d4bc89b7c915a6d80a1efbee3a708d58e688c81 (last modified on Thu Mar 26 01:11:34 2026).
WARNING:datasets.packaged_modules.cache.cache:Found the latest cached dataset configuration 'default' at /home/ubuntu/training/datasets/indic-asr-benchmark-6k/BayAreaBoys___indic-asr-benchmark-6k/default/0.0.0/4d4bc89b7c915a6d80a1efbee3a708d58e688c81 (last modified on Thu Mar 26 01:11:34 2026).
The following generation flags are not valid and may be ignored: ['top_p', 'top_k']. Set `TRANSFORMERS_VERBOSITY=info` for more details.


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check output


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


Found the latest cached dataset configuration 'default' at /home/ubuntu/training/datasets/indic-asr-benchmark-6k/BayAreaBoys___indic-asr-benchmark-6k/default/0.0.0/4d4bc89b7c915a6d80a1efbee3a708d58e688c81 (last modified on Thu Mar 26 01:11:34 2026).
WARNING:datasets.packaged_modules.cache.cache:Found the latest cached dataset configuration 'default' at /home/ubuntu/training/datasets/indic-asr-benchmark-6k/BayAreaBoys___indic-asr-benchmark-6k/default/0.0.0/4d4bc89b7c915a6d80a1efbee3a708d58e688c81 (last modified on Thu Mar 26 01:11:34 2026).
The following generation flags are not valid and may be ignored: ['top_p', 'top_k']. Set `TRANSFORMERS_VERBOSITY=info` for more details.


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


ps aux | grep 115734 | grep -v grep 2>&1 && nvidia-smi --query-compute-apps=pid,used_gpu_memory --format=csv 2>&1

Check process and GPU


ubuntu    115734  143  1.8 36889188 2295788 ?    RLl  01:18   1:11 python3 benchmark_gemma3n_asr.py --checkpoint /home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-10000 --checkpoint-name ckpt-10000 --output-dir /home/ubuntu/training/benchmark_results_gemma3n --batch-size 1
pid, used_gpu_memory [MiB]
115734, 18722 MiB


wc -l /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Count output lines


16 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check latest output


Running inference (batch_size=1)...
  [10/6000] (0%) processed


Good, it's running — just 10 samples so far but the output is buffered. The process is using 143% CPU and 18GB GPU. Let me wait for it to make more progress.


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check latest output


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check latest output


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check latest output


Loading dataset: BayAreaBoys/indic-asr-benchmark-6k
Dataset: 6000 samples

Running inference (batch_size=1)...
  [10/6000] (0%) processed


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check output


Dataset: 6000 samples

Running inference (batch_size=1)...
  [10/6000] (0%) processed
  [20/6000] (0%) processed


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check output


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check output


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check output


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check output


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check output


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check output


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check output


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


Running inference (batch_size=1)...
  [10/6000] (0%) processed
  [20/6000] (0%) processed
  [30/6000] (0%) processed


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


Running inference (batch_size=1)...
  [10/6000] (0%) processed
  [20/6000] (0%) processed
  [30/6000] (0%) processed
  [40/6000] (1%) processed


The benchmark is running at ~10 samples/min. At this rate, the full 6000-sample run will take about 10 hours. Here's what's done so far:

**Status:**
- ckpt-10000: Downloaded, benchmark running (30/6000 processed)
- ckpt-20000: Downloaded (model files only, ~21GB)
- Benchmark script `/home/ubuntu/training/benchmark_gemma3n_asr.py` working with:
  - Tokenizer patching for transformers version compatibility
  - KV cache enabled for faster generation
  - SDPA attention
  - Audio token fix (`<audio_soft_token>` not `<audio>`)

**Verified:** The model produces meaningful transcriptions — tested on Hindi samples with near-perfect output.

The main bottleneck is sequential processing (batch_size=1) since the processor handles one audio at a time. I'll keep monitoring and when ckpt-10000 finishes, I'll immediately start ckpt-20000.


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check benchmark progress


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check benchmark progress


tail -5 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check benchmark progress


  [10/6000] (0%) processed
  [20/6000] (0%) processed
  [30/6000] (0%) processed
  [40/6000] (1%) processed
  [50/6000] (1%) processed


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


  [30/6000] (0%) processed
  [40/6000] (1%) processed
  [50/6000] (1%) processed


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


  [40/6000] (1%) processed
  [50/6000] (1%) processed
  [60/6000] (1%) processed


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


can we do offline batching for faster ? and why are we not using FA2.


tail -3 /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output 2>&1

Check progress


Good questions.

**FA2:** Flash Attention 2 isn't installed on this machine — `flash_attn` package is missing. I'm using SDPA (scaled dot product attention) as a fallback, which is built into PyTorch.

**Offline batching:** Yes, we can speed this up significantly. The current bottleneck is sequential per-sample processing. Let me install FA2 and implement batched inference.


Install flash attention 2


Command running in background with ID: bdq6svc4i. Output is being written to: /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/bdq6svc4i.output


While FA2 installs, let me kill the slow run and implement offline batched inference. The key idea: pre-decode all audio, then process in batches with padding.


kill 115734 2>&1; echo "Killed slow benchmark"

Kill slow sequential benchmark


Killed slow benchmark


<task-notification>
<task-id>brrrsmz06</task-id>
<tool-use-id>toolu_01YX2njcBNARqBxf6r7piPEB</tool-use-id>
<output-file>/tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output</output-file>
<status>failed</status>
<summary>Background command "Full benchmark ckpt-10000 with KV cache + SDPA" failed with exit code 143</summary>
</task-notification>
Read the output file to retrieve the result: /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/brrrsmz06.output


<task-notification>
<task-id>bdq6svc4i</task-id>
<tool-use-id>toolu_01UQm5Vg7NxwnuY3cvoBryuc</tool-use-id>
<output-file>/tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/bdq6svc4i.output</output-file>
<status>completed</status>
<summary>Background command "Install flash attention 2" completed (exit code 0)</summary>
</task-notification>
Read the output file to retrieve the result: /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/bdq6svc4i.output


Now let me rewrite the inference to use offline batching. The challenge is that the Gemma3nProcessor doesn't natively support batched audio+text. I'll pre-process all audio features and then batch the generation.


    99→    def __call__(
   100→        self,
   101→        images: Optional[ImageInput] = None,
   102→        text: Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]] = None,
   103→        audio: Optional[Union[np.ndarray, list[float], list[np.ndarray], list[list[float]]]] = None,
   104→        videos=None,
   105→        **kwargs: Unpack[Gemma3nProcessorKwargs],
   106→    ) -> BatchFeature:
   107→        if text is None and images is None and audio is None:
   108→            raise ValueError("Provide at least one of `text`, `images`, or `audio`.")
   109→
   110→        output_kwargs = self._merge_kwargs(
   111→            Gemma3nProcessorKwargs,
   112→            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
   113→            **kwargs,
   114→        )
   115→
   116→        if isinstance(text, str):
   117→            text = [text]
   118→        elif not isinstance(text, list) and not isinstance(text[0], str):
   119→            raise ValueError("Invalid input text. Please provide a string, or a list of strings")
   120→
   121→        if audio is not None:
   122→            audio_inputs = self.feature_extractor(audio, **output_kwargs["audio_kwargs"])
   123→
   124→            if not text:
   125→                text = [self.audio_token for _ in audio]
   126→
   127→            # Expand placeholder audio tokens to the full audio token sequence
   128→            text = [prompt.replace(self.audio_token, self.full_audio_sequence) for prompt in text]
   129→        else:
   130→            audio_inputs = {}
   131→
   132→        if images is not None:
   133→            images = self.image_processor.fetch_images(images)
   134→            batched_images = make_nested_list_of_images(images)
   135→            image_inputs = self.image_processor(batched_images, **output_kwargs["images_kwargs"])
   136→
   137→            # Create empty text to be replaced with placeholders
   138→            if not text:
   139→                text = [" ".join([self.image_token] * len(images)) for images in batched_images]
   140→
   141→            if len(batched_images) != len(text):
   142→                raise ValueError(
   143→                    f"Received inconsistently sized batches of images ({len(batched_images)}) and text ({len(text)})."
   144→                )
   145→
   146→            # Expand placeholder image tokens to the full image token sequence
   147→            text = [prompt.replace(self.image_token, self.full_image_sequence) for prompt in text]
   148→        else:
   149→            image_inputs = {}
   150→
   151→        return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
   152→        text_inputs = self.tokenizer(text=text, **output_kwargs["text_kwargs"], return_tensors="np")
   153→        self._check_special_mm_tokens(text, text_inputs, modalities=["image"])
   154→
   155→        # Add token type ids manually, as tokenizer can't do arbitrary position token types
   156→        array_ids = text_inputs["input_ids"]
   157→        token_type_ids = np.zeros_like(array_ids)
   158→        token_type_ids[array_ids == self.image_token_id] = 1
   159→        token_type_ids[array_ids == self.audio_token_id] = 3
   160→        text_inputs = {k: v.tolist() for k, v in text_inputs.items()}  # in case user requested list inputs
   161→        text_inputs["token_type_ids"] = token_type_ids.tolist()
   162→        return BatchFeature(data={**text_inputs, **image_inputs, **audio_inputs}, tensor_type=return_tensors)
   163→

    def __call__(
        self,
        images: Optional[ImageInput] = None,
        text: Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]] = None,
        audio: Optional[Union[np.ndarray, list[float], list[np.ndarray], list[list[float]]]] = None,
        videos=None,
        **kwargs: Unpack[Gemma3nProcessorKwargs],
    ) -> BatchFeature:
        if text is None and images is None and audio is None:
            raise ValueError("Provide at least one of `text`, `images`, or `audio`.")

        output_kwargs = self._merge_kwargs(
            Gemma3nProcessorKwargs,
            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
            **kwargs,
        )

        if isinstance(text, str):
            text = [text]
        elif not isinstance(text, list) and not isinstance(text[0], str):
            raise ValueError("Invalid input text. Please provide a string, or a list of strings")

        if audio is not None:
            audio_inputs = self.feature_extractor(audio, **output_kwargs["audio_kwargs"])

            if not text:
                text = [self.audio_token for _ in audio]

            # Expand placeholder audio tokens to the full audio token sequence
            text = [prompt.replace(self.audio_token, self.full_audio_sequence) for prompt in text]
        else:
            audio_inputs = {}

        if images is not None:
            images = self.image_processor.fetch_images(images)
            batched_images = make_nested_list_of_images(images)
            image_inputs = self.image_processor(batched_images, **output_kwargs["images_kwargs"])

            # Create empty text to be replaced with placeholders
            if not text:
                text = [" ".join([self.image_token] * len(images)) for images in batched_images]

            if len(batched_images) != len(text):
                raise ValueError(
                    f"Received inconsistently sized batches of images ({len(batched_images)}) and text ({len(text)})."
                )

            # Expand placeholder image tokens to the full image token sequence
            text = [prompt.replace(self.image_token, self.full_image_sequence) for prompt in text]
        else:
            image_inputs = {}

        return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
        text_inputs = self.tokenizer(text=text, **output_kwargs["text_kwargs"], return_tensors="np")
        self._check_special_mm_tokens(text, text_inputs, modalities=["image"])

        # Add token type ids manually, as tokenizer can't do arbitrary position token types
        array_ids = text_inputs["input_ids"]
        token_type_ids = np.zeros_like(array_ids)
        token_type_ids[array_ids == self.image_token_id] = 1
        token_type_ids[array_ids == self.audio_token_id] = 3
        text_inputs = {k: v.tolist() for k, v in text_inputs.items()}  # in case user requested list inputs
        text_inputs["token_type_ids"] = token_type_ids.tolist()
        return BatchFeature(data={**text_inputs, **image_inputs, **audio_inputs}, tensor_type=return_tensors)


python3 -c "
import numpy as np, torch, time
from transformers import AutoTokenizer, AutoConfig, AutoProcessor, AutoModelForImageTextToText
from datasets import load_dataset

ckpt = '/home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-10000'
config = AutoConfig.from_pretrained(ckpt, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(ckpt, trust_remote_code=True)
patches = {
    'audio_token_id': config.audio_token_id, 'audio_token': '<audio_soft_token>',
    'boa_token_id': config.boa_token_id, 'boa_token': '<start_of_audio>',
    'eoa_token_id': config.eoa_token_id, 'eoa_token': '<end_of_audio>',
    'boi_token_id': config.boi_token_id, 'boi_token': '<start_of_image>',
    'eoi_token_id': config.eoi_token_id, 'eoi_token': '<end_of_image>',
    'image_token_id': config.image_token_id, 'image_token': '<image_soft_token>',
}
for k, v in patches.items():
    if not hasattr(tokenizer, k):
        setattr(tokenizer, k, v)

processor = AutoProcessor.from_pretrained(ckpt, trust_remote_code=True, tokenizer=tokenizer)
model = AutoModelForImageTextToText.from_pretrained(ckpt, dtype=torch.bfloat16, device_map='cuda:0', trust_remote_code=True, attn_implementation='sdpa')
model.eval()
model.config.use_cache = True
model.config.text_config.use_cache = True

ds = load_dataset('BayAreaBoys/indic-asr-benchmark-6k', cache_dir='/home/ubuntu/training/datasets/indic-asr-benchmark-6k', split='train')
ds = ds.filter(lambda x: x['lang_code'] == 'hi')

# Test batch of 4
batch_size = 4
wavs = []
prompts = []
for i in range(batch_size):
    s = ds[i]
    wav = np.array(s['audio']['array'], dtype=np.float32)
    sr = s['audio']['sampling_rate']
    if sr != 16000:
        import torchaudio
        wav = torchaudio.functional.resample(torch.from_numpy(wav).unsqueeze(0), sr, 16000).squeeze(0).numpy()
    wavs.append(wav)
    msg = [{'role': 'user', 'content': '<audio_soft_token>\nTranscribe the audio.\nLanguage: Hindi'}]
    prompts.append(processor.apply_chat_template(msg, add_generation_prompt=True, tokenize=False))

t0 = time.time()
inputs = processor(text=prompts, audio=wavs, sampling_rate=16000, return_tensors='pt', padding=True)
print(f'Processor time: {time.time()-t0:.2f}s')
print('input_ids shape:', inputs['input_ids'].shape)
print('input_features shape:', inputs.get('input_features', torch.tensor([])).shape)

inputs = {k: v.to(model.device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}

t0 = time.time()
with torch.no_grad():
    out = model.generate(**inputs, max_new_tokens=256, do_sample=False)
print(f'Generate time: {time.time()-t0:.2f}s for {batch_size} samples')

for i in range(batch_size):
    gen_ids = out[i][inputs['input_ids'].shape[1]:]
    text = processor.tokenizer.decode(gen_ids, skip_special_tokens=True).strip()
    ref = ds[i]['reference']
    print(f'  [{i}] REF: {ref[:50]}...')
    print(f'  [{i}] HYP: {text[:50]}...')
" 2>&1

Test batched inference