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.
Processor time: 0.04s
input_ids shape: torch.Size([4, 210])
input_features shape: torch.Size([4, 599, 128])
Generate time: 2.46s for 4 samples
  [0] REF: इस तरह इसका सोलहदिवसीय क्रम लगातार जारी है...
  [0] HYP: इस तरह इसका सोलह दिवसीय क्रम लगातार जारी है...
  [1] REF: इसके अलावा आलराउंडर कोरी एंडरसन और ग्रांट इलियट भी...
  [1] HYP: इसके अलावा ऑल राउंडर कोरी एंडरसन और ग्रांट एलियट भ...
  [2] REF: हमारी सरकार ने सिंचाई की व्यवस्था की है...
  [2] HYP: हमारी सरकार ने सिंचाई की व्यवस्था की है...
  [3] REF: पुणे और मुंबई का आंतरराष्ट्रीय हवाईअड्डा अनुक्रमे ...
  [3] HYP: पुणे और मुंबई का अंतरराष्ट्रीय हवाई अड्डा अनुक्रमे...


Batched inference works. **4 samples in 2.46s** vs ~5.7s sequential (1.43s × 4). That's 2.3x speedup. Let me try larger batch sizes.


Bash

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')

for batch_size in [8, 16, 32]:
    wavs = []
    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)

    prompts = [processor.apply_chat_template([{'role': 'user', 'content': '<audio_soft_token>\nTranscribe the audio.\nLanguage: Hindi'}], add_generation_prompt=True, tokenize=False)] * batch_size

    inputs = processor(text=prompts, audio=wavs, sampling_rate=16000, return_tensors='pt', padding=True)
    inputs = {k: v.to(model.device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}

    torch.cuda.synchronize()
    t0 = time.time()
    with torch.no_grad():
        out = model.generate(**inputs, max_new_tokens=256, do_sample=False)
    torch.cuda.synchronize()
    elapsed = time.time() - t0
    rate = batch_size / elapsed * 60
    print(f'batch_size={batch_size}: {elapsed:.2f}s ({rate:.0f} samples/min)')
" 2>&1

Test larger batch sizes


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.
batch_size=8: 2.51s (191 samples/min)
batch_size=16: 2.17s (442 samples/min)
batch_size=32: 2.41s (796 samples/min)


Massive speedup with batching:
- batch_size=8: **191 samples/min**
- batch_size=16: **442 samples/min**
- batch_size=32: **796 samples/min**

At 796 samples/min, 6000 samples = **~7.5 minutes**! Let me update the benchmark script to support batched inference.


    65→        wav_tensor = torchaudio.functional.resample(wav_tensor, sr, target_sr)
    66→        wav = wav_tensor.squeeze(0).numpy()
    67→
    68→    return wav
    69→
    70→
    71→def run_inference(model, processor, dataset, batch_size: int = 1, max_new_tokens: int = 256):
    72→    """Run inference on HF dataset and return list of result dicts."""
    73→    results = []
    74→    total = len(dataset)
    75→
    76→    for start in range(0, total, batch_size):
    77→        end = min(start + batch_size, total)
    78→        batch = dataset.select(range(start, end))
    79→
    80→        for i, sample in enumerate(batch):
    81→            lang_name = LANG_CODE_TO_NAME.get(sample["lang_code"], sample["language"].capitalize())
    82→
    83→            # Construct prompt — use <audio_soft_token> which processor expands to full audio sequence
    84→            user_content = f"<audio_soft_token>\nTranscribe the audio.\nLanguage: {lang_name}"
    85→            messages = [{"role": "user", "content": user_content}]
    86→
    87→            # Get audio array
    88→            wav = load_audio_array(sample)
    89→
    90→            try:
    91→                # Apply chat template
    92→                prompt_text = processor.apply_chat_template(
    93→                    messages, add_generation_prompt=True, tokenize=False
    94→                )
    95→
    96→                # Process inputs (text + audio)
    97→                inputs = processor(
    98→                    text=prompt_text,
    99→                    audio=[wav],
   100→                    sampling_rate=16000,
   101→                    return_tensors="pt",
   102→                )
   103→
   104→                # Move to model device
   105→                inputs = {k: v.to(model.device) if isinstance(v, torch.Tensor) else v
   106→                          for k, v in inputs.items()}
   107→
   108→                # Generate
   109→                with torch.no_grad():
   110→                    output_ids = model.generate(
   111→                        **inputs,
   112→                        max_new_tokens=max_new_tokens,
   113→                        do_sample=False,
   114→                    )
   115→
   116→                # Decode only the new tokens
   117→                input_len = inputs["input_ids"].shape[1]
   118→                generated_ids = output_ids[0][input_len:]
   119→                hypothesis = processor.tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
   120→
   121→            except Exception as e:
   122→                print(f"  Sample {start + i} ({sample['id']}) FAILED: {e}")
   123→                hypothesis = ""
   124→
   125→            results.append({
   126→                "id": sample["id"],
   127→                "language": sample["language"],
   128→                "lang_code": sample["lang_code"],
   129→                "source": sample["source"],
   130→                "duration": sample["duration"],
   131→                "reference": sample["reference"],
   132→                "hypothesis": hypothesis,
   133→            })
   134→

        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,
            })


Edit

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 run_inference(model, processor, dataset, batch_size: int = 16, max_new_tokens: int = 256):
    """Run batched inference on HF dataset and return list of result dicts."""
    results = []
    total = len(dataset)

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

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

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

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

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

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

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

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

    return results


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

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

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

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

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

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

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

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

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

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

    return results

#!/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,
        attn_implementation="sdpa",
    )
    model.eval()
    # Enable KV cache for faster autoregressive generation (training disables it)
    model.config.use_cache = True
    if hasattr(model.config, "text_config"):
        model.config.text_config.use_cache = True

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

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

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

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

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

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

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

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

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

    # Print table
    print_metrics_table(metrics, ckpt_name)

    # Save sample analysis
    save_sample_analysis(results, output_dir)

if __name__ == "__main__":
    main()

-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."""

+def run_inference(model, processor, dataset, batch_size: int = 16, max_new_tokens: int = 256):

+    """Run batched inference on HF dataset and return list of result dicts."""

-        for i, sample in enumerate(batch):

+        # Prepare batch: audio arrays + prompts

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

-            wav = load_audio_array(sample)

-                prompt_text = processor.apply_chat_template(

-                    messages, add_generation_prompt=True, tokenize=False

-                # Process inputs (text + audio)

-                    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()}

-                    output_ids = model.generate(

-                        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()

-                print(f"  Sample {start + i} ({sample['id']}) FAILED: {e}")

+            wavs.append(load_audio_array(sample))

+            prompts.append(processor.apply_chat_template(

+                messages, add_generation_prompt=True, tokenize=False

-                "hypothesis": hypothesis,

+                text=prompts, audio=wavs, sampling_rate=16000,

+                return_tensors="pt", padding=True,

+            inputs = {k: v.to(model.device) if isinstance(v, torch.Tensor) else v

+                      for k, v in inputs.items()}

+            input_len = inputs["input_ids"].shape[1]

+                output_ids = model.generate(

+                    **inputs, max_new_tokens=max_new_tokens, do_sample=False,

+            # Decode each sample in the batch

+            for i in range(len(batch_meta)):

+                gen_ids = output_ids[i][input_len:]

+                hypothesis = processor.tokenizer.decode(gen_ids, skip_special_tokens=True).strip()

+                results.append({**batch_meta[i], "hypothesis": hypothesis})

+            print(f"  Batch {start}-{end} FAILED ({e}), falling back to sequential...", flush=True)

+            for i, sample in enumerate(batch):

+                        text=[prompts[i]], audio=[wavs[i]], sampling_rate=16000,

+                        return_tensors="pt",

+                    inp = {k: v.to(model.device) if isinstance(v, torch.Tensor) else v

+                           for k, v in inp.items()}

+                    with torch.no_grad():

+                        out = model.generate(**inp, max_new_tokens=max_new_tokens, do_sample=False)

+                    gen_ids = out[0][inp["input_ids"].shape[1]:]

+                except Exception as e2:

+                    print(f"  Sample {start+i} ({batch_meta[i]['id']}) FAILED: {e2}")

-        if end % max(batch_size * 10, 10) == 0 or end == total:

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

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


Also update the default batch_size in the argparse.


    parser.add_argument("--batch-size", type=int, default=1,
                        help="Samples per batch (1 for sequential inference)")

    parser.add_argument("--batch-size", type=int, default=16,
                        help="Samples per batch for offline batched inference")


    parser.add_argument("--batch-size", type=int, default=16,
                        help="Samples per batch for offline batched inference")

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

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

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

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

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

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

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

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

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

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

    return wav

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

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

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

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

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

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

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

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

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

    return results

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

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

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

        if not ref_norm:
            continue

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

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

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

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

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

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

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

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

    return metrics

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

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

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

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

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

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

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

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

def main():
    parser = argparse.ArgumentParser(description="Benchmark Gemma3n-E2B-ASR on Indic ASR dataset")
    parser.add_argument("--checkpoint", required=True, help="Path to model checkpoint directory")
    parser.add_argument("--checkpoint-name", default=None,
                        help="Name for this checkpoint in output files (default: dirname)")
    parser.add_argument("--hf-dataset", default="BayAreaBoys/indic-asr-benchmark-6k",
                        help="HuggingFace dataset name")
    parser.add_argument("--hf-token", default=None, help="HuggingFace token (or set HF_TOKEN env)")
    parser.add_argument("--cache-dir", default="/home/ubuntu/training/datasets/indic-asr-benchmark-6k")
    parser.add_argument("--output-dir", default="/home/ubuntu/training/benchmark_results_gemma3n")
    parser.add_argument("--batch-size", type=int, default=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,
        attn_implementation="sdpa",
    )
    model.eval()
    # Enable KV cache for faster autoregressive generation (training disables it)
    model.config.use_cache = True
    if hasattr(model.config, "text_config"):
        model.config.text_config.use_cache = True

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

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

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

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

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

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

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

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

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

    # Print table
    print_metrics_table(metrics, ckpt_name)

    # Save sample analysis
    save_sample_analysis(results, output_dir)

if __name__ == "__main__":
    main()

-    parser.add_argument("--batch-size", type=int, default=1,

-                        help="Samples per batch (1 for sequential inference)")

+    parser.add_argument("--batch-size", type=int, default=16,

+                        help="Samples per batch for offline batched inference")


Now let me check if FA2 installed, then run the full benchmark with batching.


Bash

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

Check FA2 install progress


      ModuleNotFoundError: No module named 'torch'
      [end of output]

  note: This error originates from a subprocess, and is likely not a problem with pip.
ERROR: Failed to build 'flash-attn' when getting requirements to build wheel


FA2 build failed (torch not found during build). SDPA + batching is already very fast. Let me run the full benchmark now.


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 16 2>&1

Full benchmark ckpt-10000 with batching


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


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

Check batched benchmark progress


(Bash completed with no output)


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

Check progress


Check full output


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

Check process


ubuntu    120048  0.0  0.0   7504  3584 ?        Ss   01:24   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 16 2>&1' \< /dev/null && pwd -P >| /tmp/claude-6583-cwd
ubuntu    120068  131  1.2 12785812 1578452 ?    Rl   01:24   0:17 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 16


Read output


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

Check file


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


     1→`torch_dtype` is deprecated! Use `dtype` instead!
     2→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']
     3→- 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).
     4→- 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).
     5→Using the latest cached version of the dataset since BayAreaBoys/indic-asr-benchmark-6k couldn't be found on the Hugging Face Hub
     6→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
     7→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).
     8→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).
     9→The following generation flags are not valid and may be ignored: ['top_p', 'top_k']. Set `TRANSFORMERS_VERBOSITY=info` for more details.
    10→

`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/bu1n3uikr.output 2>&1

Check 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.


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

Check progress


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

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


Running with batches now. Let me keep checking.


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

Check progress


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

Check progress


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

Check progress


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

Check progress


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

Check progress


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

Check progress


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

Check progress


Dataset: 6000 samples

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


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

Check progress


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

Check progress


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

Check progress


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

Check progress


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

Check progress


Running inference (batch_size=16)...
  [16/6000] (0%) processed
  [32/6000] (1%) processed
  [48/6000] (1%) processed


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

Check progress


Loading checkpoint: /home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-10000
Model loaded in 18.1s
Loading dataset: BayAreaBoys/indic-asr-benchmark-6k
Dataset: 6000 samples

Running inference (batch_size=16)...
  [16/6000] (0%) processed
  [32/6000] (1%) processed
  [48/6000] (1%) processed
  [64/6000] (1%) processed


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

Check progress


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

Running inference (batch_size=16)...
  [16/6000] (0%) processed
  [32/6000] (1%) processed
  [48/6000] (1%) processed
  [64/6000] (1%) processed
  [80/6000] (1%) processed
  [96/6000] (2%) processed


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

Check progress


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

Check progress


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

Check progress


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

Check progress


  [48/6000] (1%) processed
  [64/6000] (1%) processed
  [80/6000] (1%) processed
  [96/6000] (2%) processed
  [112/6000] (2%) processed


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

Check progress


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

Check progress


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

Check progress


  [64/6000] (1%) processed
  [80/6000] (1%) processed
  [96/6000] (2%) processed
  [112/6000] (2%) processed
  [128/6000] (2%) processed


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

Check progress


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

Check progress


  [80/6000] (1%) processed
  [96/6000] (2%) processed
  [112/6000] (2%) processed
  [128/6000] (2%) processed
  [144/6000] (2%) processed


Good — about 144 samples processed. The rate seems to be around 16 samples every ~30 seconds = ~32 samples/min. Let me wait for it to finish. At this rate, ETA is about 3 hours. Let me keep monitoring periodically.


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

Check progress


  [112/6000] (2%) processed
  [128/6000] (2%) processed
  [144/6000] (2%) processed


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

Check progress


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

Check progress


  [128/6000] (2%) processed
  [144/6000] (2%) processed
  [160/6000] (3%) processed


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

Check progress


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

Check progress


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

Check progress


  [144/6000] (2%) processed
  [160/6000] (3%) processed
  [176/6000] (3%) processed


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

Check progress


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

Check progress


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

Check progress


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

Check progress


  [160/6000] (3%) processed
  [176/6000] (3%) processed
  [192/6000] (3%) processed


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

Check progress


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

Check progress


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

Check progress


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

Check progress


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

Check progress


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

Check progress


  [176/6000] (3%) processed
  [192/6000] (3%) processed
  [208/6000] (3%) processed


Benchmark running at ~32 samples/min with batching. ETA ~3 hours. I'll continue monitoring.


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

Check progress


  [192/6000] (3%) processed
  [208/6000] (3%) processed
  [224/6000] (4%) processed


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

Check progress


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

Check progress


  [208/6000] (3%) processed
  [224/6000] (4%) processed
  [240/6000] (4%) processed


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

Check progress


  [224/6000] (4%) processed
  [240/6000] (4%) processed
  [256/6000] (4%) processed


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

Check progress


  [240/6000] (4%) processed
  [256/6000] (4%) processed
  [272/6000] (5%) processed


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

Check progress


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

Check progress


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

Check progress


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

Check progress


  [256/6000] (4%) processed
  [272/6000] (5%) processed
  [288/6000] (5%) processed


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

Check progress


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

Check progress


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

Check progress


  [272/6000] (5%) processed
  [288/6000] (5%) processed
  [304/6000] (5%) processed


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

Check progress


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

Check progress


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

Check progress


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

Check progress


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

Check progress


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

Check progress


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

Check progress


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

Check progress


  [288/6000] (5%) processed
  [304/6000] (5%) processed
  [320/6000] (5%) processed


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

Check progress


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

Check progress


  [304/6000] (5%) processed
  [320/6000] (5%) processed
  [336/6000] (6%) processed


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

Check progress


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

Check progress


  [320/6000] (5%) processed
  [336/6000] (6%) processed
  [352/6000] (6%) processed


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

Check progress


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

Check progress


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

Check progress


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

Check progress


  [336/6000] (6%) processed
  [352/6000] (6%) processed
  [368/6000] (6%) processed


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

Check progress


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

Check progress


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

Check progress


  [352/6000] (6%) processed
  [368/6000] (6%) processed
  [384/6000] (6%) processed


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

Check progress


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

Check progress


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

Check progress


  [368/6000] (6%) processed
  [384/6000] (6%) processed
  [400/6000] (7%) processed


400/6000 processed (~33 samples/min). The benchmark will finish in about 2.5 more hours. It's running stably with no errors. I'll continue monitoring periodically.


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

Check progress


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

Check progress


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

Check progress


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

Check progress


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

Check progress


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

Check progress


  [384/6000] (6%) processed
  [400/6000] (7%) processed
  [416/6000] (7%) processed


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

Check progress


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

Check progress


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

Check progress


  [400/6000] (7%) processed
  [416/6000] (7%) processed
  [432/6000] (7%) processed


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

Check progress


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

Check progress


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

Check progress


  [416/6000] (7%) processed
  [432/6000] (7%) processed
  [448/6000] (7%) processed


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

Check progress


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

Check progress


  [432/6000] (7%) processed
  [448/6000] (7%) processed
  [464/6000] (8%) processed


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

Check progress


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

Check progress


  [448/6000] (7%) processed
  [464/6000] (8%) processed
  [480/6000] (8%) processed


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

Check progress


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

Check progress


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

Check progress


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

Check progress


  [480/6000] (8%) processed
  [496/6000] (8%) processed
A decoder-only architecture is being used, but right-padding was detected! For correct generation results, please set `padding_side='left'` when initializing the tokenizer.


There's a right-padding warning, but since all prompts have identical length (same format + 188 audio tokens), it shouldn't affect results. Let me fix it for future runs though.


   286→
   287→    ckpt_name = args.checkpoint_name or Path(args.checkpoint).name
   288→    output_dir = Path(args.output_dir) / ckpt_name
   289→    output_dir.mkdir(parents=True, exist_ok=True)
   290→
   291→    # Load model and processor
   292→    print(f"Loading checkpoint: {args.checkpoint}")
   293→    t0 = time.time()

    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()


Edit

tokenizer = AutoTokenizer.from_pretrained(args.checkpoint, trust_remote_code=True)

tokenizer = AutoTokenizer.from_pretrained(args.checkpoint, trust_remote_code=True, padding_side="left")


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

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

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

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

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

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

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

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

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

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

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

    return wav

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

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

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

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

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

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

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

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

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

    return results

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

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

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

        if not ref_norm:
            continue

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

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

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

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

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

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

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

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

    return metrics

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    # Print table
    print_metrics_table(metrics, ckpt_name)

    # Save sample analysis
    save_sample_analysis(results, output_dir)

if __name__ == "__main__":
    main()

-    tokenizer = AutoTokenizer.from_pretrained(args.checkpoint, trust_remote_code=True)

+    tokenizer = AutoTokenizer.from_pretrained(args.checkpoint, trust_remote_code=True, padding_side="left")


Good. Let me continue monitoring the current run.


Bash

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

Check progress


  [512/6000] (9%) processed
  [528/6000] (9%) processed
  [544/6000] (9%) processed


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

Check progress


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

Check progress


  [528/6000] (9%) processed
  [544/6000] (9%) processed
  [560/6000] (9%) processed


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

Check progress


  [560/6000] (9%) processed
  [576/6000] (10%) processed
  [592/6000] (10%) processed


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

Check progress


Good, 592/6000 (10%). Progressing steadily. No errors. Let me check less frequently now.


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

Check progress


  [592/6000] (10%) processed
  [608/6000] (10%) processed
  [624/6000] (10%) processed


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

Check progress


  [608/6000] (10%) processed
  [624/6000] (10%) processed
  [640/6000] (11%) processed


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

Check progress


  [624/6000] (10%) processed
  [640/6000] (11%) processed
  [656/6000] (11%) processed


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

Check progress


  [640/6000] (11%) processed
  [656/6000] (11%) processed
  [672/6000] (11%) processed


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

Check progress


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

Check progress


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

Check progress


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

Check progress


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

Check progress