Blame for make_mobile.py

#!/usr/bin/env python3
"""Makes the mobile copy of each multilingual Whisper size the app downloads for German: the same weights with the
attention outputs removed, so ONNX Runtime does not keep a 1500 x 1500 map per head and layer for each transcription.
It is the same change as ../whisper_en/make_mobile.py, made to Xenova's multilingual exports.

Xenova's exports declare every attention map as a graph output (encoder_attentions.N, decoder_attentions.N,
cross_attentions.N). ONNX Runtime keeps declared outputs alive, so Whisper Medium needed 4.4-4.6 GB on a 12 GB
phone and Android killed the app in the foreground. Removing the outputs changes no weights, nodes or results.

For each size this script:
  1. strips the encoder (keeps last_hidden_state) and the merged decoder (keeps logits and present.*), checks
     both with onnx.checker, and writes them to out/<size>/repo/onnx/;
  2. checks the copy against the original: the same greedy tokens, and the same encoder output, on 2 FLEURS
     clips with 30 decoder steps. A mismatch stops the script;
  3. measures the net peak memory of the original and the copy, each in a fresh process (check.peak_mb);
  4. copies the tokenizer and configs of the same Xenova commit, writes LICENSE, README.md (model card) and
     out/<size>/memory.json, and copies this script into the repo as provenance.

Downloads the pinned Xenova files on first use (check.resolve with source=True). Nothing is uploaded here:
publish.py does that, after review.

Usage (from tools/whisper_de):
  ../mobile_export/.venv/bin/python -I make_mobile.py tiny base small medium
"""

import argparse
import json
import os
import platform
import shutil
import sys
from datetime import datetime, timezone
from pathlib import Path
from string import Template

import numpy as np
import onnx
import onnxruntime as ort
import soundfile as sf
from huggingface_hub import hf_hub_download
from huggingface_hub.utils import EntryNotFoundError
from transformers import WhisperFeatureExtractor

# -I leaves this folder off sys.path, so the sibling modules are found explicitly.
sys.path.insert(0, str(Path(__file__).resolve().parent))
import check  # noqa: E402
from common import APP_LICENSES, DECODER_PATH, ENCODER_PATH, OUT, SIZES, SOURCE  # noqa: E402

CONFIG_FILES = ("tokenizer.json", "config.json", "generation_config.json", "preprocessor_config.json")
ATTENTION_PREFIXES = ("encoder_attentions.", "decoder_attentions.", "cross_attentions.")
EQUALITY_CLIPS = 2
EQUALITY_STEPS = 30
SAMPLE_RATE_HZ = 16_000
CLIP_FRAMES = 3000  # 30 s of log-mel frames, the encoder's fixed input length

CARD = Template("""---
license: apache-2.0
language:
- de
- multilingual
base_model: openai/whisper-$size
pipeline_tag: automatic-speech-recognition
tags:
- onnx
- whisper
- android
- mobile
---

# Whisper $size for mobile (ONNX)

An int8 ONNX copy of OpenAI's multilingual Whisper $size model, for on-device speech recognition with ONNX
Runtime. The Thì Thầm app uses it for German, with the decoder forced to `<|de|>`.

It is made from [Xenova/whisper-$size](https://huggingface.co/Xenova/whisper-$size) at commit `$commit`.
The only change is that the attention maps are no longer graph outputs. ONNX Runtime keeps declared graph
outputs in memory, so Xenova's files need much more memory for each transcription than the model needs.
The weights, the nodes and the results are unchanged.

## Peak memory

Net peak memory of one transcription (both sessions loaded, the encoder, then 20 decoder steps), with ONNX
Runtime $ort, 4 intra-op threads and graph optimization ENABLE_ALL, measured in fresh processes:

| Model | Xenova file | This copy |
|---|---|---|
| $size | $before MB | $after MB |

## Check against the original

On $clips FLEURS German clips, $steps greedy decoder steps from the German prompt give identical tokens, and the
encoder output is identical, for Xenova's file and for this copy.

## Inputs and outputs

- Encoder `$encoder`: input `input_features`, float32 [1, 80, 3000] (log-mel of 30 s of audio); output `last_hidden_state`.
- Decoder `$decoder`: inputs `input_ids` (int64), `encoder_hidden_states`, `use_cache_branch` (bool) and `past_key_values.*`; outputs `logits` and `present.*`.
- German prompt: `[50258, 50261, 50359, 50363]` (`<|startoftranscript|>`, `<|de|>`, `<|transcribe|>`, `<|notimestamps|>`). Stop at `50257` (`<|endoftext|>`). Another language token works the same way.

## Credits and license

- Whisper: OpenAI, [openai/whisper-$size](https://huggingface.co/openai/whisper-$size) (Apache 2.0 on Hugging Face;
  the code and weights in [openai/whisper](https://github.com/openai/whisper) are under the MIT License).
- ONNX export: Xenova, [Xenova/whisper-$size](https://huggingface.co/Xenova/whisper-$size), made for transformers.js.
- This copy: made with `make_mobile.py` (included here for provenance; it runs from `tools/whisper_de` of the app's
  repository). License: Apache 2.0, see `LICENSE`, as the files it is made from.
- Used by the Thì Thầm Android app for German speech recognition.
""")



def strip(src: Path, dst: Path, keep) -> list[str]:
    """Writes src to dst with only the graph outputs that keep() accepts. Returns the names that were dropped.

    An If node's then and else branches must have the same outputs as the node, so they are filtered at the
    same positions.
    """
    model = onnx.load(str(src))
    dropped = [o.name for o in model.graph.output if not keep(o.name)]
    kept = [o for o in model.graph.output if keep(o.name)]
    del model.graph.output[:]
    model.graph.output.extend(kept)
    for node in model.graph.node:
        if node.op_type != "If":
            continue
        positions = [k for k, name in enumerate(node.output) if keep(name)]
        for attr in node.attribute:
            if attr.name in ("then_branch", "else_branch"):
                branch = list(attr.g.output)
                del attr.g.output[:]
                attr.g.output.extend(branch[k] for k in positions)
        outs = list(node.output)
        del node.output[:]
        node.output.extend(outs[k] for k in positions)
    onnx.checker.check_model(model)
    dst.parent.mkdir(parents=True, exist_ok=True)
    onnx.save(model, str(dst))
    return dropped


def keep_encoder(name: str) -> bool:
    return name == check.OUT_HIDDEN


def keep_decoder(name: str) -> bool:
    return name == check.OUT_LOGITS or name.startswith(check.PRESENT_PREFIX)


def features_of(wav: Path, extractor: WhisperFeatureExtractor) -> np.ndarray:
    audio, rate = sf.read(wav, dtype="float32")
    if rate != SAMPLE_RATE_HZ:
        sys.exit(f"error: {wav} is {rate} Hz")
    if audio.ndim > 1:
        audio = audio.mean(axis=1)
    return extractor(audio, sampling_rate=SAMPLE_RATE_HZ, return_tensors="np").input_features.astype(np.float32)


def run_model(encoder: Path, decoder: Path, features: np.ndarray) -> tuple[np.ndarray, list[int], list[np.ndarray]]:
    hidden = check.app_session(encoder).run([check.OUT_HIDDEN], {check.IN_FEATURES: features})[0]
    tokens, logits = check.decode_greedy(check.app_session(decoder), hidden, EQUALITY_STEPS)
    return hidden, tokens, logits


def check_equal(original: tuple[Path, Path], copy: tuple[Path, Path]) -> dict:
    """Asserts identical greedy tokens and encoder output on the first FLEURS clips. Returns the measured diffs."""
    extractor = WhisperFeatureExtractor(feature_size=check.MEL_BINS)
    worst_hidden = worst_logits = 0.0
    for wav, _, name in check.clips(EQUALITY_CLIPS):
        features = features_of(wav, extractor)
        hidden_a, tokens_a, logits_a = run_model(*original, features)
        hidden_b, tokens_b, logits_b = run_model(*copy, features)
        if tokens_a != tokens_b:
            sys.exit(f"error: {name}: greedy tokens differ between the original and the copy")
        worst_hidden = max(worst_hidden, float(np.abs(hidden_a - hidden_b).max()))
        worst_logits = max(worst_logits, max(float(np.abs(a - b).max()) for a, b in zip(logits_a, logits_b)))
    if worst_hidden != 0.0:
        sys.exit(f"error: encoder output differs by {worst_hidden} after stripping")
    return {"clips": EQUALITY_CLIPS, "steps": EQUALITY_STEPS, "tokens_identical": True,
            "max_abs_diff_last_hidden_state": worst_hidden, "max_abs_diff_logits": worst_logits}


def copy_configs(size: str, repo: Path) -> list[str]:
    """Copies the tokenizer and configs of the pinned Xenova commit. Returns the names that exist upstream."""
    pin = SOURCE[size]
    copied = []
    for name in CONFIG_FILES:
        try:
            src = hf_hub_download(pin.repo, name, revision=pin.commit)
        except EntryNotFoundError:
            print(f"  {name}: not in {pin.repo}, skipped")
            continue
        shutil.copy2(src, repo / name)
        copied.append(name)
    return copied


def write_card(size: str, repo: Path, memory: dict) -> None:
    pin = SOURCE[size]
    (repo / "README.md").write_text(
        CARD.substitute(
            size=size, commit=pin.commit, ort=ort.__version__,
            before=memory["original_peak_mb"], after=memory["copy_peak_mb"],
            clips=EQUALITY_CLIPS, steps=EQUALITY_STEPS, encoder=ENCODER_PATH, decoder=DECODER_PATH,
        ),
        encoding="utf-8",
    )


def build(size: str) -> dict:
    pin = SOURCE[size]
    repo = OUT / size / "repo"
    if repo.exists():
        shutil.rmtree(repo)
    print(f"== {size}: {pin.repo} at {pin.commit}")

    original = check.resolve(size, source=True)
    copy_paths = (repo / ENCODER_PATH, repo / DECODER_PATH)
    dropped_enc = strip(original[0], copy_paths[0], keep_encoder)
    dropped_dec = strip(original[1], copy_paths[1], keep_decoder)
    for label, dropped in (("encoder", dropped_enc), ("decoder", dropped_dec)):
        stray = [n for n in dropped if not n.startswith(ATTENTION_PREFIXES)]
        if stray:
            sys.exit(f"error: {label} would drop non-attention outputs {stray}")
        print(f"  {label}: dropped {len(dropped)} outputs" + (f", {dropped[0]} .. {dropped[-1]}" if dropped else ""))

    equality = check_equal(original, copy_paths)
    print(f"  equal: tokens identical on {EQUALITY_CLIPS} clips, last_hidden_state max diff "
          f"{equality['max_abs_diff_last_hidden_state']}, logits max diff {equality['max_abs_diff_logits']}")

    memory = {
        "size": size,
        "source_commit": pin.commit,
        "original_peak_mb": check.peak_mb(*original),
        "copy_peak_mb": check.peak_mb(*copy_paths),
    }
    print(f"  peak: original {memory['original_peak_mb']} MB, copy {memory['copy_peak_mb']} MB")

    copied = copy_configs(size, repo)
    shutil.copy2(APP_LICENSES / "apache-2.0.txt", repo / "LICENSE")
    shutil.copy2(Path(__file__).resolve(), repo / "make_mobile.py")
    write_card(size, repo, memory)

    memory.update({
        "onnxruntime": ort.__version__,
        "machine": platform.machine(),
        "intra_op_threads": min(check.APP_THREADS, os.cpu_count() or 1),
        "measured_utc": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "equality": equality,
        "configs": copied,
    })
    (OUT / size / "memory.json").write_text(json.dumps(memory, indent=2) + "\n", encoding="utf-8")
    print(f"  wrote {repo}")
    return memory


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("sizes", nargs="+", choices=SIZES)
    args = parser.parse_args()

    results = [build(size) for size in args.sizes]
    print()
    print(f"{'size':8} {'Xenova MB':>10} {'copy MB':>9}")
    for r in results:
        print(f"{r['size']:8} {r['original_peak_mb']:>10} {r['copy_peak_mb']:>9}")


if __name__ == "__main__":
    main()