madjinn/kb-whisper-large-coreml

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

This is KBLab/kb-whisper-large, 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), 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.

Files

kb-whisper-large-int8/:

• MelSpectrogram.mlmodelc: input audio Float16 [480000]; output melspectrogram_features [1,128,1,3000].
• AudioEncoder.mlmodelc: output encoder_output_embeds [1,1280,1,1500].
• TextDecoder.mlmodelc (32 layers):
  • inputs: input_ids, cache_length, key_cache/value_cache [1,40960,1,448], kv_cache_update_mask, encoder_output_embeds and decoder_key_padding_mask;
  • outputs: logits [1,1,51866], key_cache_updates, value_cache_updates and alignment_heads_weights.
• config.json, generation_config.json and tokenizer.json: KBLab's files at the source revision, unchanged.

18 files, 1,561,164,738 bytes (encoder 639 MB, decoder 910 MB).

How it was made

• Source: KBLab/kb-whisper-large at revision d5d5984b4d8f7c4847a8ea203f1976285fb28300.
• 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-large --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 142 dB PSNR, decoder 133 dB with the same top token everywhere.
• Transcription, greedy, Swedish, with the compiled models on a Mac's CPU, driven the way WhisperKit drives them, against an fp16 conversion made the same way:
  • three Swedish clips (8, 10 and 15 s): identical to fp16, for example "Kom hit med jollen Olsson, kommer direktörn och snart skiljer sig jollens siluett från jaktens.";
  • the encoder output's cosine similarity to fp16 is 0.9991–0.9993;
  • with an English clip and the English language token, fp16 answered in Swedish (a translation) and 8-bit in English. KB-Whisper is for Swedish.
• Not yet measured: speed and memory on devices. At 1.56 GB it is about the size of fp16 Medium models, which don't run well on every iPhone's Neural Engine.

Use with WhisperKit

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