example.py

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