From 83f205e819f8cc3034e6977dc6eff447d9646be5 Mon Sep 17 00:00:00 2001
From: gsdali <51393997+gsdali@users.noreply.github.com>
Date: Fri, 29 May 2026 09:42:37 +1000
Subject: [PATCH] Disable token timestamps for distilled models detected by
 config, not name
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The existing guard only disables token-level timestamps for models whose
version string contains 'distil'. Some distilled models are not named 'distil-*'
but still ship the original (e.g. large-v3) generation_config, whose
alignment_heads reference decoder layers that no longer exist after
distillation. Generating token timestamps then crashes with an IndexError in
compute_alignment_heads_attention_weights.

Detect this case from the config — any alignment_heads entry referencing a
layer >= decoder_layers — and disable token timestamps, in addition to the
existing name-based check.

Repro: kotoba-tech/kotoba-whisper-bilingual-v1.0 (2 decoder layers, but ships
large-v3 alignment_heads up to layer 25) fails the WhisperTextDecoder step
without this.
---
 scripts/generate_model.py | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/scripts/generate_model.py b/scripts/generate_model.py
index 8c90862..fa39413 100644
--- a/scripts/generate_model.py
+++ b/scripts/generate_model.py
@@ -133,6 +133,29 @@ def cli():
             "Disabling token-level timestamps due to missing alignment_heads in distil-whisper-* models"
         )
         args.disable_token_timestamps = True
+    else:
+        # Some distilled models are not named "distil-*" yet still carry the
+        # original (e.g. large-v3) generation_config, whose alignment_heads
+        # reference decoder layers that no longer exist after distillation.
+        # Configuring token-level timestamps then crashes with an IndexError in
+        # compute_alignment_heads_attention_weights. Detect this from the config
+        # (rather than the model name) and disable token timestamps.
+        try:
+            from transformers import GenerationConfig, WhisperConfig
+            decoder_layers = WhisperConfig.from_pretrained(
+                args.model_version).decoder_layers
+            alignment_heads = getattr(
+                GenerationConfig.from_pretrained(args.model_version),
+                "alignment_heads", None,
+            ) or []
+            if any(layer >= decoder_layers for layer, _head in alignment_heads):
+                logger.info(
+                    "Disabling token-level timestamps: alignment_heads reference "
+                    f"decoder layers beyond decoder_layers={decoder_layers}"
+                )
+                args.disable_token_timestamps = True
+        except Exception as e:
+            logger.warning(f"alignment_heads validation skipped: {e}")
 
     # Generate WhisperTextDecoder
     args.test_seq_len = args.text_decoder_max_sequence_length