"""Minimal inference example for Whisper Small INT8 using ExecuTorch. Loads a quantized .pte model and transcribes a single audio file. The INT8 model was exported via Optimum-ExecuTorch with 8da8w quantization on all Linear layers plus a manual weight-only INT8 pass on `decoder.embed_tokens`, and uses separate 'encoder' and 'text_decoder' ExecuTorch methods with a static KV cache. This example resolves the model, tokenizer, and preprocessor artifacts from the repository root. Original baseline artifacts are retained in the pte_original directory and are not used by this optimized example. Requirements: pip install executorch torch transformers soundfile numpy """ import argparse import json import time from pathlib import Path import numpy as np import torch from executorch.runtime import Runtime from transformers import AutoTokenizer # -- Configuration ------------------------------------------------------------- AUDIO_PATH = "sample_input.flac" PREPROCESSOR_FILENAME = "whisper_preprocessor.pte" MODEL_FILENAME = "whisper_small_vivo_executorch_optimized.pte" DECODER_START_TOKEN_ID = 50258 FORCED_PREFIX_IDS = [50259, 50359, 50363] # <|en|>, <|transcribe|>, <|notimestamps|> EOS_TOKEN_ID = 50257 MAX_GENERATION_TOKENS = 128 MAX_SECONDS_PER_SAMPLE = 120.0 REPETITION_GUARD_REPEATS = 3 REPETITION_GUARD_MIN_PATTERN_LEN = 2 REPETITION_GUARD_MAX_PATTERN_LEN = 16 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, 357, 366, 438, 532, 685, 705, 796, 930, 1058, 1220, 1267, 1279, 1303, 1343, 1377, 1391, 1635, 1782, 1875, 2162, 2361, 2488, 3467, 4008, 4211, 4600, 4808, 5299, 5855, 6329, 7203, 9609, 9959, 10563, 10786, 11420, 11709, 11907, 13163, 13697, 13700, 14808, 15306, 16410, 16791, 17992, 19203, 19510, 20724, 22305, 22935, 27007, 30109, 30420, 33409, 34949, 40283, 40493, 40549, 47282, 49146, 50359, 50360, 50361, ) BEGIN_SUPPRESS_TOKENS = (220, 50257) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Run Whisper Small ExecuTorch inference from an exported model bundle." ) parser.add_argument( "--model-dir", default=None, help=( "Directory containing the ExecuTorch model, tokenizer files, and optionally " "whisper_preprocessor.pte. Defaults to the directory containing example.py." ), ) parser.add_argument( "--audio", default=AUDIO_PATH, help="Path to the input audio file (.flac/.wav).", ) return parser.parse_args() def resolve_model_dir(script_dir: Path, requested_dir: str | None) -> Path: candidates: list[Path] = [] if requested_dir: candidates.append(Path(requested_dir)) candidates.append(script_dir) for candidate in candidates: bundle_dir = candidate.resolve() if (bundle_dir / MODEL_FILENAME).exists(): return bundle_dir searched = "\n".join(f"- {candidate.resolve()}" for candidate in candidates) raise FileNotFoundError( "Could not find a Whisper Small ExecuTorch model bundle. Searched:\n" f"{searched}" ) def load_audio(audio_path: str) -> tuple[np.ndarray, int]: import soundfile as sf waveform, sample_rate = sf.read(str(audio_path), dtype="float32") if waveform.ndim == 2: waveform = waveform.mean(axis=1) return waveform, int(sample_rate) def resample_to_16k(waveform: np.ndarray, sample_rate: int) -> np.ndarray: if sample_rate == 16000: return waveform target_len = int(round(len(waveform) * 16000 / sample_rate)) resampled = np.interp( np.linspace(0, len(waveform) - 1, target_len), np.arange(len(waveform)), waveform, ) return resampled.astype(np.float32) def load_preprocessor(preprocessor_path: Path): if not preprocessor_path.exists(): return None, None runtime = Runtime.get() program = runtime.load_program(str(preprocessor_path)) method_names = sorted(program.method_names) method_name = "forward" if "forward" in method_names else method_names[0] return program, program.load_method(method_name) def preprocess(audio_path: str, preprocessor_method) -> torch.Tensor: waveform, sample_rate = load_audio(audio_path) waveform_16k = resample_to_16k(waveform, sample_rate) waveform_tensor = torch.from_numpy(waveform_16k).float().contiguous() if preprocessor_method is not None: outputs = preprocessor_method.execute([waveform_tensor]) features = outputs[0] if isinstance(features, (list, tuple)): features = features[0] return torch.as_tensor(features).float().contiguous() from executorch.extension.audio.mel_spectrogram import WhisperAudioProcessor fallback_preprocessor = WhisperAudioProcessor( feature_size=80, max_audio_len=300, stack_output=True, ) with torch.no_grad(): features = fallback_preprocessor(waveform_tensor) return features.float().contiguous() def load_model(pte_path: str) -> tuple: runtime = Runtime.get() program = runtime.load_program(pte_path) available = sorted(program.method_names) print(f" Available methods: {available}") if "encoder" in available and "text_decoder" in available: return program, { "format": "seq2seq", "encoder": program.load_method("encoder"), "decoder": program.load_method("text_decoder"), } if "forward" in available: return program, { "format": "forward", "forward": program.load_method("forward"), } raise RuntimeError(f"Unknown export format. Methods found: {available}") def apply_suppression(scores: torch.Tensor, first_free_step: bool) -> torch.Tensor: vocab_size = scores.shape[-1] out = scores.clone() valid_suppress = [ t for t in SUPPRESS_TOKENS if 0 <= t < vocab_size and t != EOS_TOKEN_ID ] if valid_suppress: out[0, valid_suppress] = float("-inf") if first_free_step: valid_begin = [t for t in BEGIN_SUPPRESS_TOKENS if 0 <= t < vocab_size] if valid_begin: out[0, valid_begin] = float("-inf") return out def decode_tokens(tokenizer, token_ids: list[int]) -> str: return tokenizer.decode( token_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False, ).strip() def find_repeated_suffix_pattern( token_ids: list[int], *, repeats: int = REPETITION_GUARD_REPEATS, min_pattern_len: int = REPETITION_GUARD_MIN_PATTERN_LEN, max_pattern_len: int = REPETITION_GUARD_MAX_PATTERN_LEN, ) -> int | None: total = len(token_ids) upper = min(max_pattern_len, total // repeats) for pattern_len in range(min_pattern_len, upper + 1): pattern = token_ids[-pattern_len:] if all( token_ids[-pattern_len * (idx + 1) : -pattern_len * idx or None] == pattern for idx in range(repeats) ): return pattern_len return None def transcribe_seq2seq( encoder_method, decoder_method, features: torch.Tensor, tokenizer, ) -> dict: encoder_outputs = encoder_method.execute([features]) encoder_hidden = encoder_outputs[0] if isinstance(encoder_hidden, (list, tuple)): encoder_hidden = encoder_hidden[0] encoder_hidden = torch.as_tensor(encoder_hidden).float().contiguous() forced_prefix = list(FORCED_PREFIX_IDS) tokens = [DECODER_START_TOKEN_ID] cache_position = 0 forced_prefix_idx = 0 generated_token_count = 0 stop_reason = "max_tokens" started = time.perf_counter() generated_free_tokens: list[int] = [] for _step in range(MAX_GENERATION_TOKENS + len(forced_prefix)): input_tensor = torch.tensor([[tokens[-1]]], dtype=torch.long).contiguous() pos_tensor = torch.tensor([cache_position], dtype=torch.long).contiguous() decoder_outputs = decoder_method.execute([input_tensor, encoder_hidden, pos_tensor]) flat_logits = decoder_outputs[0] if isinstance(flat_logits, (list, tuple)): flat_logits = flat_logits[0] flat_logits = torch.as_tensor(flat_logits).float().flatten() if forced_prefix_idx < len(forced_prefix): next_token = forced_prefix[forced_prefix_idx] forced_prefix_idx += 1 else: scores = flat_logits.unsqueeze(0) first_free = generated_token_count == 0 scores = apply_suppression(scores, first_free_step=first_free) next_token = int(scores[0].argmax().item()) generated_token_count += 1 generated_free_tokens.append(next_token) if next_token == EOS_TOKEN_ID: stop_reason = "eos" tokens.append(next_token) cache_position += 1 break repeated_suffix_len = find_repeated_suffix_pattern(generated_free_tokens) if repeated_suffix_len is not None: trim_count = repeated_suffix_len * REPETITION_GUARD_REPEATS del generated_free_tokens[-trim_count:] del tokens[-(trim_count - 1) :] generated_token_count -= trim_count stop_reason = "repetition_guard" break if time.perf_counter() - started >= MAX_SECONDS_PER_SAMPLE: stop_reason = "timeout" break tokens.append(next_token) cache_position += 1 elapsed = time.perf_counter() - started text = decode_tokens(tokenizer, tokens) return { "transcription": text, "generated_tokens": generated_token_count, "stop_reason": stop_reason, "elapsed_s": round(elapsed, 3), } def transcribe_forward(forward_method, features: torch.Tensor, tokenizer) -> dict: prompt = [DECODER_START_TOKEN_ID] + list(FORCED_PREFIX_IDS) decoder_ids = torch.tensor([prompt], dtype=torch.long).contiguous() generated_token_count = 0 stop_reason = "max_tokens" started = time.perf_counter() generated_free_tokens: list[int] = [] with torch.no_grad(): for _step in range(MAX_GENERATION_TOKENS): outputs = forward_method.execute([features, decoder_ids]) logits = outputs[0] if isinstance(logits, (list, tuple)): logits = logits[0] logits = torch.as_tensor(logits).float() next_token_scores = logits[:, -1, :] first_free = generated_token_count == 0 next_token_scores = apply_suppression(next_token_scores, first_free_step=first_free) next_token = next_token_scores.argmax(dim=-1, keepdim=True).long() decoder_ids = torch.cat([decoder_ids, next_token], dim=1) generated_token_count += 1 generated_free_tokens.append(int(next_token.item())) if EOS_TOKEN_ID >= 0 and bool(torch.all(next_token == EOS_TOKEN_ID)): stop_reason = "eos" break repeated_suffix_len = find_repeated_suffix_pattern(generated_free_tokens) if repeated_suffix_len is not None: trim_count = repeated_suffix_len * REPETITION_GUARD_REPEATS generated_free_tokens = generated_free_tokens[:-trim_count] decoder_ids = decoder_ids[:, :-trim_count] generated_token_count -= trim_count stop_reason = "repetition_guard" break if time.perf_counter() - started >= MAX_SECONDS_PER_SAMPLE: stop_reason = "timeout" break elapsed = time.perf_counter() - started text = decode_tokens(tokenizer, decoder_ids[0].tolist()) return { "transcription": text, "generated_tokens": generated_token_count, "stop_reason": stop_reason, "elapsed_s": round(elapsed, 3), } def save_results(result: dict, script_dir: Path) -> None: output_path = script_dir / "transcription.json" with open(output_path, "w", encoding="utf-8") as f: json.dump(result, f, indent=2, ensure_ascii=False) print(f"Saved transcription to {output_path}") def main() -> None: args = parse_args() script_dir = Path(__file__).parent model_dir = resolve_model_dir(script_dir, args.model_dir) model_path = model_dir / MODEL_FILENAME audio_arg = Path(args.audio) audio_path = audio_arg if audio_arg.is_absolute() else (script_dir / audio_arg).resolve() preprocessor_path = model_dir / PREPROCESSOR_FILENAME print(f"Loading tokenizer from {model_dir} ...") tokenizer = AutoTokenizer.from_pretrained( str(model_dir), local_files_only=True, use_fast=True, ) _preprocessor_program = None preprocessor_method = None if preprocessor_path.exists(): print(f"Loading preprocessor from {preprocessor_path} ...") _preprocessor_program, preprocessor_method = load_preprocessor(preprocessor_path) else: print("Local preprocessor .pte not found; falling back to WhisperAudioProcessor.") print(f"Loading model from {model_path} ...") program, methods = load_model(str(model_path)) print(f"Preprocessing audio: {audio_path}") features = preprocess(str(audio_path), preprocessor_method) print(f" Input features shape: {tuple(features.shape)}") print("Running transcription ...") if methods["format"] == "seq2seq": result = transcribe_seq2seq( methods["encoder"], methods["decoder"], features, tokenizer ) else: result = transcribe_forward(methods["forward"], features, tokenizer) print(f"\nTranscription: {result['transcription']!r}") print(f"Generated tokens: {result['generated_tokens']}") print(f"Stop reason: {result['stop_reason']}") print(f"Elapsed: {result['elapsed_s']:.3f} s") save_results(result, script_dir) if __name__ == "__main__": main()