diff --git a/docs/word-prediction.md b/docs/word-prediction.md index fe957c2b..6bba6c28 100644 --- a/docs/word-prediction.md +++ b/docs/word-prediction.md @@ -14,7 +14,7 @@ One Word prediction setting controls an offline SmolLM2-135M int8 ONNX model. Th For a first typed prefix, a fixed local context keeps the model from favoring website names at the start of a document; it adds no user text. The model loads in the worker while keyboard input remains available. Suggestions are blank until loading finishes. A passive badge beside the scan prompt distinguishes loading, a ready keyboard awaiting typed context, no matching suggestions, available suggestions, paused activity tracking, and prediction failure. It never shows typed text and is not a scan target. If loading or inference fails, or a call exceeds 1.5 seconds, the keyboard continues accepting input and offers **Retry predictions** in its toolbar. Retry restarts only the prediction worker, clears its private text context, and leaves the keyboard open; type a new prefix afterward. A 400 ms search budget bounds candidate exploration. A clipped buffer with no complete earlier word yields no suggestions. -Candidates contain ASCII letters and apostrophes and must have sufficient model probability. There is no vocabulary filter: any word the model finds likely can be suggested, including swearing, because the person typing chose it. Single-letter candidates are limited to “a” and “I”. Up to five suggestions are shown: the first, third and fifth are the most likely single words, and the second and fourth are the two most likely two-word phrases that begin with one of the top three words, ranked by the probability of the pair. When fewer phrases are found within an extra 200 ms, single words fill the remaining slots, and the other way round. Accepting a phrase inserts the rest of its first word, a space, the second word and a trailing space. Ordinary words use lowercase; mixed-case names retain model casing. Casing is applied when the keyboard displays a candidate, and the label shows exactly what will be inserted. Caps, or Shift locked, uppercases the whole completion, and together they cancel as they do on typed letters. Before any letter of the word is typed, Shift once changes only the first letter, the way it would change the next typed letter, and a sentence start capitalises it without any modifier. Shift once under Caps gives a lowercase first letter there, as it would on a typed letter. After typed letters, a pending Shift once is left for the next letter and does not touch the completion. +Candidates contain ASCII letters and apostrophes and must have sufficient model probability. There is no vocabulary filter: any word the model finds likely can be suggested, including swearing, because the person typing chose it. Single-letter candidates are limited to “a” and “I”. Up to five suggestions are shown: the first, third and fifth are the most likely single words, and the second and fourth are the two most likely two-word phrases that begin with one of the top three words, ranked by the probability of the pair. When fewer phrases are found within an extra 200 ms, single words fill the remaining slots, and the other way round. Accepting a phrase inserts the rest of its first word, a space, the second word and a trailing space. Ordinary words use lowercase. Names keep the model's capital: mixed case anywhere, such as WhatsApp, or an initial capital mid-sentence that the model clearly prefers, such as London or Monday, which it does for proper nouns and occasionally a rare word. At a sentence start the model capitalises every word, so the lowercase form wins there and the sentence rule decides. Casing is applied when the keyboard displays a candidate, and the label shows exactly what will be inserted. Caps, or Shift locked, uppercases the whole completion, and together they cancel as they do on typed letters. Before any letter of the word is typed, Shift once changes only the first letter, the way it would change the next typed letter, and a sentence start capitalises it without any modifier. Shift once under Caps gives a lowercase first letter there, as it would on a typed letter. After typed letters, a pending Shift once is left for the next letter and does not touch the completion. The model files are too large to commit. `npm run prediction-model` downloads them from a pinned upstream revision and verifies their SHA-256. The Tauri dev and build commands and CI run it automatically. Run it once before `cargo test` or `cargo clippy` in a fresh checkout. diff --git a/src-tauri/src/prediction/model.rs b/src-tauri/src/prediction/model.rs index ccf5b6d3..6030f5d2 100644 --- a/src-tauri/src/prediction/model.rs +++ b/src-tauri/src/prediction/model.rs @@ -76,9 +76,11 @@ pub struct Prediction { } /// A word (or phrase) the beam finished: its probability mass over spellings, -/// the best spelling, that spelling's pieces and their log-probability. +/// how much of it spells the last word with a capital, the best spelling, +/// that spelling's pieces and their log-probability. struct Finished { mass: f32, + capital: f32, best: f32, spelling: String, pieces: Vec, @@ -165,14 +167,23 @@ fn grow( .map(|(l, _)| l.exp()) .sum(); let score = logp + boundary.max(1e-12).ln(); + let capital = word + .rsplit(' ') + .next() + .and_then(|last| last.chars().next()) + .is_some_and(char::is_uppercase); let entry = finished.entry(key).or_insert(Finished { mass: 0.0, + capital: 0.0, best: f32::NEG_INFINITY, spelling: String::new(), pieces: Vec::new(), logp: f32::NEG_INFINITY, }); entry.mass += score.exp(); + if capital { + entry.capital += score.exp(); + } if score > entry.best { entry.best = score; entry.spelling = word.clone(); @@ -207,16 +218,26 @@ fn ranked(finished: HashMap, limit: usize) -> Vec<(String, Fin entries } -/// How a finished word is shown. Sentence position and all-caps tokens -/// should not force uppercase onto ordinary completions; mixed-case names -/// keep the model's casing, the pronoun I is always capital, and the typed -/// apostrophe style is kept. -fn display(key: &str, spelling: &str, typed: &str) -> String { +/// Capitalised spellings must carry this share of a word's probability for +/// the capital to be a name rather than a coin toss on a rare word. +const NAME_SHARE: f32 = 0.9; + +/// How a finished word is shown. All-caps tokens never force uppercase onto +/// a completion. A name keeps the model's capital: mixed case anywhere, or an +/// initial capital mid-sentence that the model clearly prefers, which it only +/// does for a proper noun. At a sentence start the model capitalises every +/// word, so the lowercase form wins there and the sentence rule decides. The +/// pronoun I is always capital, and the typed apostrophe style is kept. +fn display(key: &str, f: &Finished, typed: &str, sentence_start: bool) -> String { + let spelling = &f.spelling; + let all_upper = spelling.chars().all(|c| c.is_ascii_uppercase()); + let mixed = spelling.chars().skip(1).any(|c| c.is_ascii_uppercase()); + let name = spelling.chars().next().is_some_and(char::is_uppercase) + && !sentence_start + && f.capital >= NAME_SHARE * f.mass; let word = if key == "i" { "I".to_owned() - } else if spelling.chars().skip(1).any(|c| c.is_ascii_uppercase()) - && !spelling.chars().all(|c| c.is_ascii_uppercase()) - { + } else if !all_upper && (mixed || name) { spelling.replace('’', "'") } else { key.to_owned() @@ -275,20 +296,24 @@ pub fn complete( } let words = ranked(grow(scorer, vocabulary, beam, &prefix, deadline)?, 5); let phrases = follow(scorer, vocabulary, &words, phrase_deadline)?; + let sentence_start = start || text.ends_with(['.', '!', '?']); Ok(Prediction { words: words .iter() - .map(|(key, f)| display(key, &f.spelling, typed)) + .map(|(key, f)| display(key, f, typed, sentence_start)) .collect(), phrases: phrases .into_iter() .map(|(lead, key, f)| { - let second = f.spelling.split_once(' ').map_or("", |(_, s)| s); let second_key = key.split_once(' ').map_or("", |(_, s)| s); + let second = Finished { + spelling: f.spelling.split_once(' ').map_or("", |(_, s)| s).to_owned(), + ..f + }; format!( "{} {}", - display(&words[lead].0, &words[lead].1.spelling, typed), - display(second_key, second, typed) + display(&words[lead].0, &words[lead].1, typed, sentence_start), + display(second_key, &second, typed, false) ) }) .collect(), @@ -716,11 +741,18 @@ mod tests { } #[test] - fn casing_comes_from_model_except_sentence_position() { + fn capital_spellings_survive_mid_sentence_but_not_at_a_sentence_start() { let mut s = fake(&[(1, 0.9)]); let lowered = ["water", "waffle", "wa"]; - assert_eq!(run(&mut s, "Send ", "").unwrap(), lowered); + assert_eq!(run(&mut s, "Send ", "").unwrap(), ["Water", "Waffle", "Wa"]); assert_eq!(run(&mut s, "Done. ", "").unwrap(), lowered); + // The capital must carry nine tenths of the word's mass: 0.8 against + // 0.1 is a share of 0.89 and stays lowercase; 0.85 against 0.05 is + // 0.94 and keeps the capital. + let mut s = fake(&[(1, 0.8), (2, 0.1)]); + assert_eq!(run(&mut s, "Send ", "").unwrap()[0], "water"); + let mut s = fake(&[(1, 0.85), (2, 0.05)]); + assert_eq!(run(&mut s, "Send ", "").unwrap()[0], "Water"); } #[test] @@ -776,6 +808,8 @@ mod tests { ("I would like a cup of ", "", Some("tea")), ("Please put the ", "ket", Some("kettle")), ("Can you send a ", "wh", None), + ("I live in ", "lon", Some("London")), + ("See you on ", "mon", Some("Monday")), ("I want to ", "fa", None), ] { let t = Instant::now();