Kol cleanup 360M (Core ML)

On-device English dictation cleanup for Kol, a macOS dictation app. The model takes a raw speech-recognition transcript plus context from the screen (vocabulary, conversation title, the focused text field, the user's custom spelling rules) and returns the same words with punctuation, capitalization, filler removal and context spellings fixed. It does not rephrase.

It is SmolLM2-360M-Instruct with a rank-16 LoRA merged in, converted to a Core ML ML Program (fp16, macOS 15+) with a stateful key-value cache.

Files

Inputs and outputs

• input_ids int32 [1, q]: the next q tokens (1 to 1,024).
• cache_position int32 [1]: where those tokens start in the cache.
• causal_mask float16 [1, 1, q, 1024]: 0 where a query may attend, -10000 elsewhere.
• State key_cache, value_cache: the 1,024-token cache.
• Output token_ids int32 [1, q]: the greedy next token after each input position. Kol uses this to verify up to 32 drafted tokens per call, copied ahead from the transcript, so most of a cleanup takes a few calls.

The chat has Kol's system prompt (compact-en-v2, in KolCore/LocalCleanupPrompt.swift) and one user message: a JSON object with text, category, vocabulary, conversation_name, focused, ide, custom_rules, mentions_enabled and language. The reply is the cleaned text, or EMPTY.

Training

LoRA fine-tune (rank 16, 5,500 steps, batch 8, learning rate 2e-4) starting from an adapter trained on Kol's fictional generated cleanup data. The mix:

• 20,000 pairs sampled from three public dictation-cleanup datasets, kept only when the cleaned text uses the speaker's own words and does not delete swear words, emphasis words (literally, honestly, seriously, totally, actually, basically) or negations: juanquivilla/sotto-transcript-cleanup (MIT), shantanugoel/aawaaz-transcript-cleanup-dataset (MIT), MagicNoThief/handy-dictation-editing (CC-BY-4.0).
• About 20,800 generated examples with fictional names covering screen vocabulary, conversation titles, code separators and operators, mentions, numbers and literal text.

No user dictations were used for training.

Evaluation

• 1,000-case fictional development suite: 89.4% exact, proper nouns 98.6%; 391/400 on the earlier suite.
• Held-out public sets (exact match): handy eval 305/400, sotto validation 185/397. These references delete emphasis words that this model keeps on purpose.
• Replayed through Kol on an M5 Pro: median cleanup 250 ms including Kol's name matching.

Limitations

English only. Up to 1,024 tokens of prompt and output together. It sometimes still drops or misspells a word ("erred" written "errred"). Kol checks every output against the transcript and pastes the transcript when a context name appears where the speech does not support it.