DEV Community

Alex Day
Alex Day

Posted on Edited on

Building next-word prediction for a keyboard app - pipeline, smoothing algorithms, and vocabulary tradeoffs

This is a summary of three recent posts from my CodeKeyboard rebuild series. CodeKeyboard is an open-source Android keyboard for programmers — layers, modifier keys, fuzzy autocorrect, and now a real statistical next-word prediction engine.

The prediction model

The keyboard trains a trigram language model on the Coursera SwiftKey corpus: 583MB of blogs, news, and Twitter text, 4.27 million lines. A trigram model stores, for every two-word context ("I am"), a ranked list of likely next words ("going", "a", "sure"...). At typing time, the keyboard looks up whatever the last two committed words were and surfaces the top candidates.

Three algorithm variants were built:

Kneser-Ney (KN) - instead of counting how often a word appears, it counts how many different contexts a word appears in. The word "Francisco" appears frequently in the corpus, but almost always after "San". KN scores it low as a generic next-word guess because it continues so few contexts. Good for cold-start guessing; harder to tune under size constraints.

Katz backoff - uses Good-Turing discounting to redistribute probability mass toward words the model has seen less of. More stable when the model gets pruned aggressively to fit on a phone.

SwiftKey WDP - takes the Katz model and asks of every stored next-word candidate: "does knowing the two-word context actually change this word's odds compared to the one-word context alone?" Candidates that merely restate the simpler model get dropped. Result: 31% smaller file, 99.2% top-1 agreement with Katz. This one ships.

The pipeline problem

The first full-corpus AWS run consumed all memory and produced nothing. The second attempt ran for over two hours and was killed manually.

The fix: switched to a streaming SQLite pipeline — counts flow through a rolling merge rather than loading into RAM, intermediate results checkpoint to disk, and the whole build now completes in 11 minutes on a spot t3.xlarge.

The vocabulary question

The trained model has a 427,651-word vocabulary. 85% of those words appear fewer than 15 times across 85 million tokens. Zipf's law in action.

Four vocabulary caps were tested: 16K, 32K, 64K, 128K words. The measurement: how often does the right next word fall outside the cap? At 64K, coverage reaches 99.3% of real next-word targets. Above that,
gains are marginal. The 64K cap ships.

The result

Everything — vocabulary, character trie, bigram/trigram follower lists, phrase data — comy .cklm file. 22MB. Loaded via memory-map at startup. Three separate JSON files are gone.APK: 302MB down to 94MB.

Second Improvement

Full write-ups:

Top comments (5)

Collapse
 
alexshev profile image
Alex Shev

Keyboard prediction is a great place for boring engineering choices to matter. Vocabulary size, smoothing, and latency probably affect user trust more than a clever model that occasionally feels slow or strange.

Collapse
 
silvern47 profile image
Alex Day

I am trying 4-gram and phrase completion. Its dang costly (time, aws bills, cpu fan speed, time) to even build a map with say the legacy Swiftkey corpus + the 2018's opensubtitles en->en dataset.
Maybe I should just try using karpatty's nanogpt or whatever, and keep the 3-gram.

Collapse
 
alexshev profile image
Alex Shev

The right baseline depends on the product constraint, but I would measure the 3-gram before replacing it. A compact model with a well-bounded vocabulary and a clear latency budget can be more useful than a larger system whose training and serving cost overwhelms the interaction.

Thread Thread
 
silvern47 profile image
Alex Day

dev.to/silvern47/minor-improvement... , I found a way to do better suggestion, using layout based proximity based weights.

Yeah the previous 3-gram was based of the decade old swiftkey corpus. Its holding up alright, I am trying out the opensubs for 3-gram.
There is also this infi-gram using suffix array that I landed on, which will explore later. You are right on the product constraint. I initially started it to be a split keyboard, but I am also thinking, for the corpus, maybe I could provided different datasets, which would be specific to normal conversation vs code-related. For code-related I would probably need to use some code specific datasets, which even these llm guys are using, prune it, and let users toggle between.

Thread Thread
 
alexshev profile image
Alex Shev

That split makes sense. Conversation and code are different vocabularies, but I would treat them as explicit modes with their own evaluation set rather than one mixed corpus. The useful signal is not only top-1 accuracy: log accepted suggestions and backspaces per mode, since a model can improve offline while feeling worse at the keyboard.