simplewords-dictation-cleanup-v2

Turns raw voice dictation into clean, punctuated written text, on-device on Apple Silicon via MLX. Invoked per utterance, where per-call latency matters more than throughput.

"um so like can you uh send the report to the team by tomorrow" → "Can you send the report to the team by tomorrow?"

This repo ships an 8-bit fused model (load it directly, no base needed) plus the LoRA adapter under adapter/. It is a QLoRA fine-tune of mlx-community/Qwen3.5-2B-MLX-4bit with LoRA on all linear layers.

What changed from v1

v1 had a disqualifying bug: self-correction inversion. It kept the replaced side of a spoken correction.

input	v1	v2
let's meet on tuesday wait no friday at noon	❌ "Let's meet on tuesday"	✅ "Let's meet on Friday at noon."
four options wait no five options	❌ kept "four"	✅ "Five options."
three thirty no wait four o'clock	❌ kept "three thirty"	✅ "Four o'clock."
red one no the blue one actually the green one	❌ "I'll take the red one"	✅ "Pick the green one."

The cause was structural, not a prompting failure — explicit rules and worked examples in the system prompt still produced the wrong side. v1's data generator was clean-anchored, so a value-changing correction could not be represented in its training set. v2 retrains on a regenerated corpus (53.5k rows, 25% self-correction with ground-truth {wrong, right, retained} labels).

v2 also fixes v1's secondary gaps: lowercase proper nouns, long inputs returned as run-ons, dropped terminal punctuation and ? on questions, over-trimming of short mostly-filler inputs, and missing list formatting.

⚠️ v1 and v2 use different SYSTEM prompts and must not be mixed. Sending v1's prompt to v2 weights (or the reverse) degrades quality. v2's prompt ships as system_v2.txt — read it from the repo, don't retype it.

Evaluation

Greedy (temperature 0). Two held-out sets: 44 hand-authored regression cases reproducing every v1 failure class verbatim, and 1,546 stratified test rows. Self-correction rows in the held-out set sit on templates that appear nowhere in training, so these measure generalization rather than memorization.

	regression (n=44)	held-out (n=1,546)
self-correction	100%	100%
overall behavioral	97.7%	99.6%
production guard pass	100%	99.9%

Per class (held-out): proper_nouns 100% · lists 100% · long_segmentation 100% · mostly_filler_short 100% · misc 100% · already_clean 99.4% · filler_stutter 99.4% · question_buried 97.8%.

Baselines on the same harness: v1 ≈ 0% self-correction; the un-tuned 4-bit base scores 18.8% (regression) / 0.0% (held-out).

Known gap: 43/44 on the regression set — one long rambling input is returned as a run-on instead of being split into sentences. Held-out long_segmentation is 100%, so this is a single hard instance rather than a class-wide weakness.

Caveat: references are synthetic (model-generated + templated), not human-transcribed real ASR. This is a strong in-distribution regression suite, not a substitute for evaluation on your own audio.

Usage

pip install mlx-lm huggingface_hub

from pathlib import Path
from huggingface_hub import snapshot_download
from mlx_lm import load, generate
from mlx_lm.sample_utils import make_sampler

path = snapshot_download("abhiram3040/simplewords-dictation-cleanup-v2")
system = (Path(path) / "system_v2.txt").read_text().strip()   # FROZEN prompt
model, tok = load(path)                                        # 8-bit fused
sampler = make_sampler(temp=0.0)                               # GREEDY

def clean(raw: str) -> str:
    prompt = tok.apply_chat_template(
        [{"role": "user", "content": f"{system}\n\n{raw}"}],
        add_generation_prompt=True,
        enable_thinking=False,      # empty pre-filled <think></think>; no trace
        tokenize=False,
    )
    return generate(model, tok, prompt=prompt, max_tokens=512, sampler=sampler).strip()

print(clean("let's meet on tuesday wait no friday at noon"))
# -> "Let's meet on Friday at noon."

A runnable version is in example.py.

Three things are load-bearing:

1. The SYSTEM prompt is frozen and folded into the user turn (there is no separate system role). It must match training byte-for-byte.
2. Decode greedily (temp=0.0).
3. enable_thinking=False. For Qwen3.5 this pre-fills an empty <think></think> block — that empty block is the mechanism that suppresses reasoning-trace leakage.

Recommended output guard

The model is trained to satisfy this guard, so a failure means something is wrong upstream (usually a prompt mismatch). Reject the output and fall back to pasting the raw text.

import re

_HEADER     = re.compile(r"(?m)^\s{0,3}#{1,6}\s")
_THINK      = re.compile(r"<think>", re.I)
_LIST_LINE  = re.compile(r"(?m)^\s*(?:-\s|\*\s|\u2022\s|\d+[.)]\s)")
# STRICT: bullets are allowed ONLY when the speaker literally said "list"/"bullet".
# The model is trained to this exact gate -- do not widen it to spoken counts
# ("three things"), or the model and the guard will disagree.
_LIST_SIGNAL = re.compile(r"\b(list|bullet|bullets)\b", re.I)
_WORD       = re.compile(r"[a-z0-9']+", re.I)


def accept_output(raw: str, out: str):
    """Return `out` if it passes every production gate, else None (paste raw instead)."""
    if out is None:
        return None
    if _HEADER.search(out) or _THINK.search(out):
        return None                                   # no markdown headers / think blocks
    if len(out) > 2 * len(raw) + 20:
        return None                                   # runaway expansion
    if len(raw) > 15 and len(out) < 0.3 * len(raw):
        return None                                   # over-trimmed
    if raw.rstrip().endswith("?") and "?" not in out:
        return None                                   # a question must stay a question
    if len(raw) <= 15:                                # short input must share a word
        if not (set(w.lower() for w in _WORD.findall(raw))
                & set(w.lower() for w in _WORD.findall(out))):
            return None
    if _LIST_LINE.search(out) and not _LIST_SIGNAL.search(raw):
        return None                                   # unrequested list formatting
    return out

⚠️ Do not re-fuse this adapter at 4-bit

Merging the adapter into the 4-bit base and re-quantizing to 4 bits rounds the low-rank delta away. Measured: held-out self-correction collapses 100% → 57.8% (129 failures vs 6) — while val loss and filler-removal scores still look perfect, so it fails silently. A self-correction is a sharp keep-this-span/drop-that-span decision and is far more weight-sensitive than filler removal.

The published 8-bit model was produced by dequantizing, then re-quantizing to 8 bits:

mlx_lm fuse --model <4bit-base> --adapter-path adapter/ --save-path fused-bf16 --dequantize
mlx_lm convert --hf-path fused-bf16 --mlx-path fused-q8 -q --q-bits 8 --q-group-size 64

Its failure set is identical to the adapter's, and it is also the fastest option — a runtime adapter adds per-layer overhead a fused model doesn't have:

artifact	size	ms/utterance	tok/s	self-correction
4-bit base + adapter	1.0 GB + 64 MB	311	27.3	100%
8-bit fused (this repo)	1.9 GB	214	39.5	100%
bf16 fused	3.5 GB	245	34.7	100%
4-bit fused	1.0 GB	—	—	57.8% ❌

Measured on an M3 Ultra Mac Studio over the 44 regression utterances.

Using the adapter instead

from mlx_lm import load
model, tok = load("mlx-community/Qwen3.5-2B-MLX-4bit",
                  adapter_path=snapshot_download(REPO) + "/adapter")

Training

base	mlx-community/Qwen3.5-2B-MLX-4bit (4-bit, QLoRA)
LoRA	rank 16, scale 2.0, dropout 0.05, all linear layers
data	48,339 train / 2,543 val, template-disjoint splits
objective	completions-only (mask_prompt)
schedule	6,000 iters (1 epoch), batch 8, grad checkpointing, lr 1e-4 cosine, warmup 150
seq len	768 (longest example 612 tokens)
hardware	M3 Ultra Mac Studio, 512 GB — 12.2 h, 36 GB peak
final val loss	0.005

The shipping checkpoint was chosen by behavioral score, not val loss — val loss flatlines from ~iter 1000 while behavior stays non-monotonic (iter 1500 passes the release gate, iter 2000 fails it, later ones pass). All 12 checkpoints were scored on the full behavioral suite; the gate was self-correction ≥95% on both sets, with val loss used only to break ties.

Intended use and limitations

Intended: post-processing short English dictation (messages, notes, to-dos, calendar, search queries, doc dictation) before it is pasted into an app.

Not intended: it is not a chat assistant. It never answers questions or follows instructions found inside the dictated text — it only cleans them up. Non-English input is untested. Very long inputs (beyond a few hundred tokens) are outside the training distribution. It will not fix ASR word errors it cannot infer from context.

License

Apache-2.0, inheriting the Qwen3.5-2B base license.

Adapter

The LoRA adapter this model was fused from lives at abhiram3040/simplewords-dictation-cleanup-v2-adapter. It is kept out of this repo so that snapshot_download here fetches only the fused model — a loader that merges every .safetensors under the snapshot directory would otherwise try to apply LoRA tensors to an already-fused graph.