"""Minimal inference example for Whisper Medium (LiteRT / TFLite, INT8). Loads the multi-signature ``whisper-medium-int8-litert.tflite`` model (encode + decode signatures with self-attention and cross-attention KV caches) and transcribes a single audio file with greedy decoding. """ import json from pathlib import Path import numpy as np import soundfile as sf from ai_edge_litert.interpreter import Interpreter from transformers import WhisperProcessor # ── Configuration ────────────────────────────────────────────────────────── SCRIPT_DIR = Path(__file__).parent MODEL_PATH = str(SCRIPT_DIR / "whisper-medium-int8-litert.tflite") AUDIO_PATH = str(SCRIPT_DIR / "sample_input.flac") PROCESSOR_DIR = str(SCRIPT_DIR) # tokenizer.json / processor_config.json live here TARGET_SAMPLE_RATE = 16000 NUM_MEL_BINS = 80 MAX_SOURCE_FRAMES = 3000 # 30s of audio at 10ms hop MAX_TARGET_POSITIONS = 448 # decoder self-attention cache length NUM_HEADS = 16 HEAD_DIM = 64 DECODER_LAYERS = 24 VOCAB_SIZE = 51865 LANGUAGE = "en" TASK = "transcribe" EOS_TOKEN_ID = 50257 MAX_GENERATION_TOKENS = 128 LARGE_NEGATIVE = -1e30 # Non-speech / special tokens masked at every decode step (Whisper-medium # generation config: timestamps, non-English language tags, etc.). SUPPRESS_TOKENS = [ 1, 2, 7, 8, 9, 10, 14, 25, 26, 27, 28, 29, 31, 58, 59, 60, 61, 62, 63, 90, 91, 92, 93, 359, 503, 522, 542, 873, 893, 902, 918, 922, 931, 1350, 1853, 1982, 2460, 2627, 3246, 3253, 3268, 3536, 3846, 3961, 4183, 4667, 6585, 6647, 7273, 9061, 9383, 10428, 10929, 11938, 12033, 12331, 12562, 13793, 14157, 14635, 15265, 15618, 16553, 16604, 18362, 18956, 20075, 21675, 22520, 26130, 26161, 26435, 28279, 29464, 31650, 32302, 32470, 36865, 42863, 47425, 49870, 50254, 50258, 50358, 50359, 50360, 50361, 50362, ] # Extra tokens masked only at the first generated step (leading space, eos). BEGIN_SUPPRESS_TOKENS = [220, 50257] # ── Model Loading ──────────────────────────────────────────────────────── def load_model(tflite_path: str): """Load the LiteRT model and return its encode/decode signature runners.""" interpreter = Interpreter(model_path=tflite_path) interpreter.allocate_tensors() return interpreter.get_signature_runner("encode"), interpreter.get_signature_runner( "decode" ) def load_processor(processor_dir: str) -> WhisperProcessor: """Load the feature extractor + tokenizer used to train/export the model.""" return WhisperProcessor.from_pretrained(processor_dir) def classify_decode_io(decode_runner) -> tuple[dict[str, str], dict[str, str]]: """Map logical tensor roles to the generic ``args_N`` / ``output_N`` names. The exported decode signature does not carry semantic tensor names, so roles are recovered from shape and dtype: - (1, 1) int64 -> decoder_token - (1,) int64 -> input_pos - (1, 1, 1, MAX_TARGET_POSITIONS) float32 -> self_attn_mask - (1, MAX_TARGET_POSITIONS, heads, dim) -> self_kv_i (small seq dim) - (1, MAX_SOURCE_POSITIONS, heads, dim) -> cross_kv_i (large seq dim) Outputs: one (1, 1, vocab) logits tensor + 2*layers updated self-KV tensors. """ in_details = decode_runner.get_input_details() sorted_in = sorted(in_details, key=lambda n: int(n.split("_")[1])) input_map: dict[str, str] = {} self_names, cross_names = [], [] for name in sorted_in: shape = tuple(int(x) for x in in_details[name]["shape"]) dtype = in_details[name]["dtype"] if dtype == np.int64 and shape == (1, 1): input_map["decoder_token"] = name elif dtype == np.int64 and shape == (1,): input_map["input_pos"] = name elif dtype == np.float32 and len(shape) == 4 and shape[:3] == (1, 1, 1): input_map["self_attn_mask"] = name already = set(input_map.values()) for name in sorted_in: if name in already: continue shape = tuple(int(x) for x in in_details[name]["shape"]) if in_details[name]["dtype"] != np.float32 or len(shape) != 4: continue (self_names if shape[1] <= 1024 else cross_names).append(name) for i, name in enumerate(self_names): input_map[f"self_kv_{i}"] = name for i, name in enumerate(cross_names): input_map[f"cross_kv_{i}"] = name out_details = decode_runner.get_output_details() sorted_out = sorted(out_details, key=lambda n: int(n.split("_")[1])) output_map: dict[str, str] = {} cache_out = [] for name in sorted_out: shape = tuple(int(x) for x in out_details[name]["shape"]) if len(shape) == 3 and shape[0] == 1 and shape[1] == 1: output_map["logits"] = name else: cache_out.append(name) for i, name in enumerate(cache_out): output_map[f"self_kv_{i}"] = name return input_map, output_map # ── Preprocessing ──────────────────────────────────────────────────────── def resample_to_16k(waveform: np.ndarray, sample_rate: int) -> np.ndarray: """Linearly resample mono audio to the model's expected 16kHz rate.""" if sample_rate == TARGET_SAMPLE_RATE or waveform.size == 0: return waveform duration_s = waveform.shape[0] / sample_rate new_length = int(round(duration_s * TARGET_SAMPLE_RATE)) old_x = np.linspace(0.0, 1.0, num=waveform.shape[0], endpoint=False) new_x = np.linspace(0.0, 1.0, num=new_length, endpoint=False) return np.interp(new_x, old_x, waveform).astype(np.float32) def preprocess(audio_path: str, processor: WhisperProcessor) -> np.ndarray: """Load an audio file and compute the (1, 80, 3000) log-mel spectrogram.""" waveform, sample_rate = sf.read(audio_path, dtype="float32") if waveform.ndim == 2: # downmix stereo to mono waveform = waveform.mean(axis=1) waveform = resample_to_16k(waveform, sample_rate) features = processor.feature_extractor( waveform, sampling_rate=TARGET_SAMPLE_RATE, return_tensors="np" ).input_features return features.astype(np.float32) # (1, 80, 3000) # ── Inference ──────────────────────────────────────────────────────────── def run_inference( encode_runner, decode_runner, input_map: dict[str, str], output_map: dict[str, str], input_features: np.ndarray, prompt_ids: list[int], ) -> list[int]: """Encode audio once, then greedily decode tokens using self/cross KV caches.""" # Encode: audio -> per-layer cross-attention K/V projections. encode_outputs = encode_runner(args_0=input_features) sorted_keys = sorted(encode_outputs, key=lambda n: int(n.split("_")[1])) cross_kv_buffers = [encode_outputs[k] for k in sorted_keys] # k_0, v_0, k_1, v_1, ... # Zero-initialized self-attention KV cache (updated every decode step). self_cache_shape = (1, MAX_TARGET_POSITIONS, NUM_HEADS, HEAD_DIM) n_self_cache = DECODER_LAYERS * 2 self_kv_buffers = [ np.zeros(self_cache_shape, dtype=np.float32) for _ in range(n_self_cache) ] generated_ids = list(prompt_ids) token = generated_ids[0] prompt_length = len(generated_ids) max_positions = min(MAX_GENERATION_TOKENS + prompt_length, MAX_TARGET_POSITIONS) for t in range(max_positions): token_arr = np.array([[token]], dtype=np.int64) pos_arr = np.array([t], dtype=np.int64) mask = np.zeros((1, 1, 1, MAX_TARGET_POSITIONS), dtype=np.float32) if t + 1 < MAX_TARGET_POSITIONS: mask[0, 0, 0, t + 1 :] = LARGE_NEGATIVE decode_inputs = { input_map["decoder_token"]: token_arr, input_map["input_pos"]: pos_arr, input_map["self_attn_mask"]: mask, } for i, buf in enumerate(self_kv_buffers): decode_inputs[input_map[f"self_kv_{i}"]] = buf for i, buf in enumerate(cross_kv_buffers): decode_inputs[input_map[f"cross_kv_{i}"]] = buf outputs = decode_runner(**decode_inputs) for i in range(n_self_cache): self_kv_buffers[i] = outputs[output_map[f"self_kv_{i}"]] logits = outputs[output_map["logits"]][0, 0].astype(np.float32) if t < prompt_length - 1: token = generated_ids[t + 1] continue step = t - (prompt_length - 1) if SUPPRESS_TOKENS: logits[SUPPRESS_TOKENS] = LARGE_NEGATIVE if step == 0 and BEGIN_SUPPRESS_TOKENS: logits[BEGIN_SUPPRESS_TOKENS] = LARGE_NEGATIVE next_token = int(np.argmax(logits)) if next_token == EOS_TOKEN_ID: break generated_ids.append(next_token) token = next_token return generated_ids # ── Postprocessing ─────────────────────────────────────────────────────── def postprocess(token_ids: list[int], processor: WhisperProcessor) -> str: """Decode generated token ids into a plain-text transcript.""" return processor.tokenizer.decode(token_ids, skip_special_tokens=True).strip() # ── Save Results ───────────────────────────────────────────────────────── def save_results(audio_path: str, transcript: str) -> None: """Save the transcript as JSON next to the script.""" result = { "audio_file": Path(audio_path).name, "transcription": transcript, "sample_rate": TARGET_SAMPLE_RATE, } out_path = SCRIPT_DIR / "transcription.json" out_path.write_text(json.dumps(result, indent=2)) print(f"Saved transcript to {out_path}") # ── Main ───────────────────────────────────────────────────────────────── def main(): processor = load_processor(PROCESSOR_DIR) encode_runner, decode_runner = load_model(MODEL_PATH) input_map, output_map = classify_decode_io(decode_runner) input_features = preprocess(AUDIO_PATH, processor) prompt_ids = [processor.tokenizer.convert_tokens_to_ids("<|startoftranscript|>")] forced_ids = processor.get_decoder_prompt_ids(language=LANGUAGE, task=TASK) prompt_ids.extend(token_id for _, token_id in sorted(forced_ids)) token_ids = run_inference( encode_runner, decode_runner, input_map, output_map, input_features, prompt_ids ) transcript = postprocess(token_ids, processor) print(f"Transcript: {transcript}") save_results(AUDIO_PATH, transcript) if __name__ == "__main__": main()