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.

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.

Use with WhisperKit

The folder follows WhisperKit's layout, so pointing WhisperKit's modelRepo at this repository and its model at kb-whisper-medium 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; nothing was retrained.