.wmlm training and evaluation

The context model the keyboard downloads per language is a small word-level transformer that reads the sentence before the caret and says what comes next. Suggestions describes what it does in the app; this page covers the training and evaluation tools in tools/wmlm/. That directory’s README.md is the command reference; the steps below are the end-to-end path.

Prerequisites: Python 3.10+, numpy and torch; distill.py also needs transformers. Run from the repo root. en_sentences.tsv comes from the corpus step in tools/wmlm/README.md (or an existing prepared TSV).

Prepare the arrays

prep.py turns a sentence TSV into the train.npz, heldout.npz, heldout.tsv and vocab.txt a run needs. Reserving piece slots lets the model spell a word the vocabulary does not list instead of emitting <unk>:

python tools/wmlm/prep.py en_sentences.tsv en_lm vocab=24576 pieces=2048 ctx=40 \
    seed_words=app/dictionaries-src/en.txt

Distill a teacher

distill.py runs a causal language model over a sample of the training sentences and maps its top-k next-word distribution onto the student’s vocabulary. The result is a soft-target file that train.py distill= reads:

python tools/wmlm/distill.py en_lm en_soft.npz teacher=Qwen/Qwen2.5-1.5B sentences=100000 topk=24

Add masks and alignment

soft_mask.py derives per-row data from the student’s tokens and the teacher’s soft targets. --emit-align writes the word-aware position map, --emit-keep writes a per-row KD/CE gate, and --weight cand writes a candidate-count weight. build_kd_weight.py turns a mask file into a three-level weight (empty 0, piece <1, word 1):

python tools/wmlm/soft_mask.py build en_lm en_soft.npz masks \
    --weight cand --emit-keep kd --emit-align
python tools/wmlm/build_kd_weight.py masks.masks.npz masks.w03.npz piece_weight=0.3

Train and evaluate

train.py keeps the plain cross-entropy target where the teacher said nothing, mixes the distillation loss on the rows the teacher filled, and evaluates up to 20,000 held-out sentences:

python tools/wmlm/train.py en_lm en_run steps=16000 batch=512 d=256 layers=6 ff=768 ctx=40 \
    distill=en_soft.npz distill_alpha=0.3 distill_align=masks.align.npz \
    distill_mask=masks.keep.kd.npz distill_weight=masks.w03.npz

For a quick screen, slice_data.py cuts a hermetic slice of the data dir and its soft targets. average_checkpoints.py averages saved checkpoints before export; evaluate the soup on heldout first.

Compare two runs

The typing harness writes a stable per-item dump with ITEMS=items.tsv; two runs’ dumps go into paired_eval.py, which reports McNemar and a bootstrap CI. Holm is opt-in: pass the metric family’s p-values to --holm (the tool evaluates one --metric per run).

cd native/wm-engine && cargo build --release --example typing_eval
ITEMS=items_a.tsv PACK=pack MODES=base,lm \
    ./target/release/examples/typing_eval ../../app/dictionaries-src/en.txt \
    ../../en_lm/heldout.tsv ../../en_context.wmlm
ITEMS=items_b.tsv PACK=pack MODES=base,lm \
    ./target/release/examples/typing_eval ../../app/dictionaries-src/en.txt \
    ../../en_lm/heldout.tsv ../../en_context_v2.wmlm
cd ../.. && python tools/wmlm/paired_eval.py \
    native/wm-engine/items_a.tsv native/wm-engine/items_b.tsv --metric lm_strip_top1

LANG=de selects QWERTZ and turns off english_sources, but the repo ships only en_bigrams.txt; gunzip the data repo’s de_bigrams.txt.gz and de_trigrams.txt.gz into a PACK directory for a German run. SEEDS is only read when english_sources=true, so German runs have no seed pairs. LANG is normalized from locales such as de_DE.UTF-8 and de-DE; other locales fall back to English. The glide harness (native/wm-engine/examples/glide_lm_blend_eval.rs) writes the same kind of per-item dump for a stroke set. Its records.jsonl schema is word, sentence, word_idx, canvas_width, canvas_height, and data (an array of {x, y, t} points normalized to the canvas).

Export

export.py writes the int8 .wmlm and a .sha256 sidecar next to it:

python tools/wmlm/export.py en_run/best.pt en_lm/vocab.txt en_context.wmlm \
    --langs=en --name=English
sha256sum -c en_context.wmlm.sha256

The app-side reader and the catalog entry are described in Suggestions; the Rust reader is native/wm-engine/src/lm.rs.