madjinn/whisper-medium-coreml

Whisper Medium for Core ML, 8-bit (WhisperKit format)

This is OpenAI's multilingual Whisper Medium, converted to Core ML for the Apple Neural Engine in the format of WhisperKit (Argmax), with its weights quantised to 8 bits. It is made for the dictation keyboard WhispAway, with exactly the recipe of our KB-Whisper Medium 8-bit (kb-whisper-medium-int8/), so the two share one graph and differ only in their weights.

Files

whisper-medium-int8/:

• MelSpectrogram.mlmodelc: input audio Float16 [480000]; output melspectrogram_features [1,80,1,3000].
• AudioEncoder.mlmodelc: output encoder_output_embeds [1,1024,1,1500].
• TextDecoder.mlmodelc:
  • inputs: input_ids, cache_length, key_cache/value_cache [1,24576,1,448], kv_cache_update_mask, encoder_output_embeds and decoder_key_padding_mask;
  • outputs: logits [1,1,51865], key_cache_updates, value_cache_updates and alignment_heads_weights.
• config.json, generation_config.json and tokenizer.json: OpenAI's files at the source revision, unchanged.

18 files, 775,151,649 bytes (encoder 308 MB, decoder 459 MB).

How it was made

• Source: openai/whisper-medium at revision abdf7c39ab9d0397620ccaea8974cc764cd0953e.
• Converter: whisperkittools at b5c1e50eb611240025c158784ec34bbfce26de09, with argmaxtools 0.1.7, coremltools 7.1, torch 2.2.1 and transformers 4.39.1. This is Argmax's 2024 toolchain.
• Command: whisperkit-generate-model --model-version openai/whisper-medium --text-decoder-max-sequence-length 448, otherwise at its defaults: encoder attention SplitHeadsQ, decoder Cat. The conversion ran on Linux, where Core ML can't run, so the Core ML prediction steps of the tool's tests were replaced by saving the .mlpackage; the PyTorch parity tests ran unchanged.
• 8 bits: coremltools 7.1 linear_quantize_weights, mode linear_symmetric (int8, per channel), on every weight of 2048 elements or more (weight_threshold=2048), in AudioEncoder and TextDecoder. The mel spectrogram stays fp16.
• Compiled with xcrun coremlcompiler (coremlc 3520.5.1).

Checks

• In PyTorch, the converted modules match Hugging Face's model: encoder 146 dB PSNR, decoder 131 dB with the same top token everywhere.
• The interface and op graph (names, types, shapes, op counts) are identical to our kb-whisper-medium-int8/; the compiled model.mil files are byte-identical to it, so only the weights differ. The interface is also Argmax's openai_whisper-medium's.
• Transcription, greedy, with the compiled models on a Mac's CPU, driven the way WhisperKit drives them, against an fp16 conversion made the same way:
  • English, the JFK sentence (11 s): identical, and exact ("And so, my fellow Americans, ask not what your country can do for you, ask what you can do for your country.");
  • Chinese, four clips from the FLEURS dev set (7–18 s): identical to fp16, and to the reference apart from names and homophones that fp16 gets wrong too (科东 for 科隆, 家南 for 迦南, 成 for 称);
  • Swedish, one clip (8 s): 2 of 16 words differ from fp16 ("skiljer" becomes "skilder"); for Swedish, use KB-Whisper;
  • the encoder output's cosine similarity to fp16 is 0.9984–0.9988.
• Not yet measured: speed on devices.

Use with WhisperKit

The folder follows WhisperKit's layout, so pointing WhisperKit's modelRepo at this repository and its model at whisper-medium-int8 should work. We haven't tested that: WhispAway runs these files with its own Core ML pipeline.

Licence and credits

• Licence: MIT, as OpenAI released Whisper's code and model weights (openai/whisper; see LICENSE). The Hugging Face copy of the source model is tagged Apache-2.0.
• Whisper: © OpenAI.
• Converter: whisperkittools and WhisperKit by Argmax (MIT).
• This repository: a format conversion by Madjinn. The weights are OpenAI's, converted to Core ML and quantised to 8 bits; nothing was retrained.