#!/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//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//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()