Blame for example.py

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