madjinn/kb-whisper-medium-coreml

KB-Whisper Medium for Core ML (WhisperKit format)

This is KBLab/kb-whisper-medium, the Swedish Whisper model from KBLab at the National Library of Sweden, converted to Core ML for the Apple Neural Engine in the format of WhisperKit (Argmax). It is made for the dictation keyboard WhispAway, as is our KB-Whisper Small conversion.

Files

kb-whisper-medium/:

• 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: KBLab's files, unchanged.

All weights are fp16, with no quantisation. The three models take about 1.5 GB in all.

kb-whisper-medium-int8/ is an 8-bit experiment; see below.

How it was made

The same recipe as our KB-Whisper Small, with a 448-token decoder cache like Argmax's openai_whisper-medium:

• Source: KBLab/kb-whisper-medium at revision 0abe10b9d7f75d0902656e5c06c5c4d549604dc5.
• 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 KBLab/kb-whisper-medium --text-decoder-max-sequence-length 448, otherwise at its defaults: encoder attention SplitHeadsQ, decoder Cat.
• Compiled with xcrun coremlcompiler (coremlc 3520.5.1).
• Compared with Argmax's openai_whisper-medium:
  • the interface (names, types, shapes) is identical;
  • the mel spectrogram is identical op for op;
  • the decoder is equivalent: coremltools 7.1 writes each LayerNorm as a 7-op chain, where Argmax's newer tools write one layer_norm;
  • the encoder uses the 2024 recipe's split-heads attention, where Argmax's 2025 Medium uses plain matmul attention;
  • the graphs are Argmax's 2024 openai_whisper-small graphs with Medium's layers and heads.
• Checked:
  • in PyTorch, the converted modules match Hugging Face's model (encoder 144 dB PSNR, decoder 119 dB with the same top token everywhere);
  • the compiled models, run with Core ML on a Mac's CPU and driven the way WhisperKit drives them (greedy, Swedish), transcribe four Swedish test clips (8–25 s) correctly, and better than our KB-Whisper Small on the same clips (for example "silhuett", "ofrälse", "Ta mig katten").
• Not yet measured: speed on devices.

8-bit variant: kb-whisper-medium-int8/ (experimental)

The same three models with the encoder's and decoder's weights quantised to 8 bits, made the same way as our KB-Whisper Small int8:

• Made from: a fresh conversion with the recipe above (same source revision, tools and command). Its interface and op graph match kb-whisper-medium/ exactly, the mel spectrogram's weights are bit-identical, and the PyTorch checks give the same result (encoder 143 dB, decoder 118 dB).
• Method: 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. Compiled with xcrun coremlcompiler (coremlc 3520.5.1).
• Unchanged: the mel spectrogram stays fp16. The interface (names, types, shapes) and the three JSON files are the same as in kb-whisper-medium/.
• Size: about 777 MB in all (encoder 313 MB, decoder 460 MB), half of the fp16 folder.
• Accuracy on eight Swedish test clips (5–60 s, greedy decoding with Core ML on a Mac's CPU, against kb-whisper-medium/):
  • four clips (5, 8, 15 and 60 s) give identical transcripts;
  • the other four (10, 25, 30 and 30 s) differ in one or two words each: "Askengrens" becomes "Askengren", "Silverbucklan" becomes "Silverbuckla", an invented name is spelt differently, and one clip that ends mid-sentence gets one more word; there are also small punctuation differences;
  • the encoder output's cosine similarity to fp16 is 0.9987–0.9995.
• Note: the decoder's output projection has an all-zero bias. coremltools quantises it with zero scales, so it still dequantises to zeros.
• Status: made to test whether 8-bit weights run Medium faster on iPhones' Neural Engine. It hasn't been measured on a device yet.

Use with WhisperKit

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

Licence and credits

• Licence: Apache-2.0, as KB-Whisper is (see LICENSE).
• KB-Whisper: © KBLab, National Library of Sweden, fine-tuned from OpenAI Whisper (MIT).
• Converter: whisperkittools and WhisperKit by Argmax (MIT).
• This repository: a format conversion by Madjinn. The weights are KBLab's, converted to Core ML fp16 (and, in kb-whisper-medium-int8/, quantised to 8 bits); nothing was retrained.