Ask the model as each word starts, and let it choose a typo's spelling

crates/cc-app/examples/fix_eval.rs

@@ -0,0 +1,282 @@
+//! Measures what auto-fix does with typos in running text: it puts one keyboard slip into a
+//! word of each sentence, runs the word-boundary decision the app makes (with and without
+//! the language model choosing between candidates), and counts right fixes, wrong fixes
+//! and words left alone.
+//!
+//! cargo run --release -p cc-app --example fix_eval -- http://127.0.0.1:8090 <text files...>
+//!
+//! The server is a llama-server with the completion model loaded. Plain prose works as
+//! input; Harper's test texts (harper-core/tests/text/*.md) are what the thresholds were
+//! tuned on.
+
+#[allow(dead_code)]
+#[path = "../src/fixgate.rs"]
+mod fixgate;
+
+use cc_correct::{AutocorrectList, Checker, PersonalDictionary, Severity};
+use cc_engine::CompletionClient;
+use std::collections::BTreeMap;
+use std::path::Path;
+
+const LEVEL: Severity = Severity::S2;
+const CONTEXT_CHARS: usize = 300;
+const NEXT_TOKENS: u32 = 40;
+const SEED: u64 = 0x5eed_cafe;
+
+struct Rng(u64);
+
+impl Rng {
+    fn next(&mut self) -> u64 {
+        self.0 ^= self.0 << 13;
+        self.0 ^= self.0 >> 7;
+        self.0 ^= self.0 << 17;
+        self.0
+    }
+    fn below(&mut self, n: usize) -> usize {
+        (self.next() % n as u64) as usize
+    }
+}
+
+const ROWS: [(&str, f32); 3] = [("qwertyuiop", 0.0), ("asdfghjkl", 0.25), ("zxcvbnm", 0.75)];
+
+/// Keys next to `c` on a QWERTY keyboard, as cc-correct's typing model counts them.
+fn neighbours(c: char) -> Vec<char> {
+    let pos = |c: char| {
+        ROWS.iter().enumerate().find_map(|(y, (row, offset))| row.find(c).map(|x| (x as f32 + offset, y as f32)))
+    };
+    let Some((x, y)) = pos(c) else { return Vec::new() };
+    ROWS.iter()
+        .flat_map(|(row, _)| row.chars())
+        .filter(|&k| k != c && pos(k).is_some_and(|(kx, ky)| (kx - x).powi(2) + (ky - y).powi(2) <= 2.0))
+        .collect()
+}
+
+/// One slip of the given kind in `word`, or None when it does not apply.
+fn slip(word: &str, kind: &str, rng: &mut Rng) -> Option<String> {
+    let mut w: Vec<char> = word.chars().collect();
+    let n = w.len();
+    match kind {
+        "neighbour" => {
+            let i = rng.below(n);
+            let near = neighbours(w[i]);
+            if near.is_empty() {
+                return None;
+            }
+            w[i] = near[rng.below(near.len())];
+        }
+        "swap" => {
+            let i = rng.below(n - 1);
+            if w[i] == w[i + 1] {
+                return None;
+            }
+            w.swap(i, i + 1);
+        }
+        "drop" => {
+            w.remove(rng.below(n));
+        }
+        "extra" => {
+            let i = rng.below(n);
+            let near = neighbours(w[i]);
+            if near.is_empty() {
+                return None;
+            }
+            w.insert(i + rng.below(2), near[rng.below(near.len())]);
+        }
+        "double" => {
+            let i = rng.below(n);
+            w.insert(i, w[i]);
+        }
+        _ => unreachable!(),
+    }
+    Some(w.into_iter().collect())
+}
+
+/// Sentences of plain prose, each with the text before it in its paragraph.
+fn sentences(text: &str) -> Vec<(String, String)> {
+    let mut out = Vec::new();
+    for para in text.split("\n\n") {
+        let para: String = para
+            .lines()
+            .filter(|l| !l.trim_start().starts_with(['#', '|', '>', '-', '*', '`']))
+            .collect::<Vec<_>>()
+            .join(" ");
+        let para = para.replace(['_', '*', '[', ']', '(', ')'], "");
+        let mut before = String::new();
+        for s in para.split_inclusive(['.', '!', '?']) {
+            let s = s.trim();
+            let words = s.split_whitespace().count();
+            if (6..=40).contains(&words) && s.chars().all(|c| c.is_ascii()) {
+                out.push((before.clone(), s.to_string()));
+            }
+            before.push_str(s);
+            before.push(' ');
+        }
+    }
+    out
+}
+
+#[derive(Default)]
+struct Tally {
+    right: usize,
+    wrong: usize,
+    left: usize,
+}
+
+fn main() -> anyhow::Result<()> {
+    let args: Vec<String> = std::env::args().skip(1).collect();
+    let (server, files) = args.split_first().expect("usage: fix_eval <server url> <text files...>");
+    let checker = Checker::new();
+    let dict = PersonalDictionary::default();
+    let auto = AutocorrectList::load(Path::new("nonexistent-autocorrect.txt"));
+    let client = CompletionClient::new(server);
+    let rt = tokio::runtime::Builder::new_current_thread().enable_all().build()?;
+    let mut rng = Rng(SEED);
+    let kinds = ["neighbour", "swap", "drop", "extra", "double"];
+
+    let mut tables: BTreeMap<String, BTreeMap<String, Tally>> = BTreeMap::new();
+    let mut wrong_lines = Vec::new();
+    let mut samples = 0;
+    for file in files {
+        let text = std::fs::read_to_string(file)?;
+        for (para_before, sentence) in sentences(&text) {
+            let words: Vec<&str> = sentence.split(' ').collect();
+            // One target per sentence: a lowercase dictionary word of four or more letters,
+            // not the first word.
+            let targets: Vec<usize> = (1..words.len())
+                .filter(|&i| {
+                    let w = words[i];
+                    w.len() >= 4 && w.chars().all(|c| c.is_ascii_lowercase()) && checker.is_word(w, &dict)
+                })
+                .collect();
+            if targets.is_empty() {
+                continue;
+            }
+            let i = targets[rng.below(targets.len())];
+            let truth = words[i];
+            let kind = kinds[rng.below(kinds.len())];
+            let Some(typo) = slip(truth, kind, &mut rng) else { continue };
+            if typo.len() < 3 || checker.is_word(&typo, &dict) {
+                continue;
+            }
+            samples += 1;
+            let mut body = format!("{para_before}{}", words[..i].join(" "));
+            if i > 0 {
+                body.push(' ');
+            }
+            body.push_str(&typo);
+            let tail_start = body.len().saturating_sub(CONTEXT_CHARS);
+            let tail_start = (tail_start..body.len()).find(|&k| body.is_char_boundary(k)).unwrap_or(0);
+            let body = &body[tail_start..];
+            let start = body.rfind(['.', '!', '?', '\n']).map_or(0, |k| k + 1);
+            let sent = body[start..].trim_start();
+            let n_chars = sent.chars().count();
+            let before = body[..body.len() - typo.len()].trim_end_matches(' ');
+
+            let listed = auto.lookup(&typo);
+            let found = checker.check(sent, &dict);
+            let at_end = |c: &&cc_correct::Candidate| c.end == n_chars && sent.ends_with(&c.original);
+            let eligible = found
+                .iter()
+                .filter(at_end)
+                .find(|c| fixgate::eligible(&c.severity, &c.rule, &c.original, &c.replacement, &LEVEL))
+                .filter(|c| fixgate::reject(&c.original, &c.replacement, true, false, |t| checker.is_word(t, &dict)).is_none());
+            let guess = found.iter().filter(at_end).find(|c| c.rule == "spell" && c.severity == Severity::S2);
+            let needs_model = listed.is_none() && (eligible.is_some_and(|c| c.rule == "spell") || (eligible.is_none() && guess.is_some()));
+            let top = if needs_model && !before.trim().is_empty() {
+                rt.block_on(client.next_tokens(before, NEXT_TOKENS)).unwrap_or_default()
+            } else {
+                Vec::new()
+            };
+            let is_word = |w: &str| checker.is_word(w, &dict);
+            let options = checker.typo_options(&typo, &dict);
+
+            // 0.1.3: the rules, with the model's veto on guessed S0 fixes.
+            let base_fix: Option<String> = if let Some(fix) = listed.clone() {
+                Some(fix)
+            } else if let Some(c) = eligible {
+                let guessed = c.rule == "spell" && c.severity == Severity::S0 && !cc_correct::is_doubled_letter_typo(&c.original, &c.replacement);
+                if guessed && (top.is_empty() || fixgate::model_prefers(&c.original, &c.replacement, &top, is_word).is_some()) {
+                    None
+                } else {
+                    Some(c.replacement.clone())
+                }
+            } else {
+                None
+            };
+            // New: the model picks between the typo's options where it clearly expects one.
+            let decide = |picked: Option<String>, doubled_too: bool| -> Option<String> {
+                if let Some(fix) = listed.clone() {
+                    Some(fix)
+                } else if let Some(c) = eligible {
+                    let guessed = c.rule == "spell" && c.severity == Severity::S0 && !cc_correct::is_doubled_letter_typo(&c.original, &c.replacement);
+                    if !guessed {
+                        // A doubled-key fix is never held back, but the model may pick another spelling.
+                        match picked {
+                            Some(p) if doubled_too && c.rule == "spell" => Some(p),
+                            _ => Some(c.replacement.clone()),
+                        }
+                    } else if top.is_empty() {
+                        None
+                    } else if picked.is_some() {
+                        picked
+                    } else if fixgate::model_prefers(&c.original, &c.replacement, &top, is_word).is_some() {
+                        None
+                    } else {
+                        Some(c.replacement.clone())
+                    }
+                } else if guess.is_some() {
+                    picked
+                } else {
+                    None
+                }
+            };
+            let pick_fix = decide(fixgate::model_pick(&typo, &options, &top), false);
+            let mut fixes = vec![
+                ("0.1.3 (rules + veto)".to_string(), base_fix.clone()),
+                ("model picks".to_string(), pick_fix.clone()),
+                ("model picks, doubled keys too".to_string(), decide(fixgate::model_pick(&typo, &options, &top), true)),
+            ];
+            for weight in [1.0, 2.0, 3.0] {
+                for typed in [true, false] {
+                    let name = format!("weight {weight}, typed word {}", if typed { "rivals" } else { "ignored" });
+                    fixes.push((name, decide(fixgate::model_pick_tuned(&typo, &options, &top, weight, typed), false)));
+                }
+            }
+            for (name, fix) in &fixes {
+                let table = tables.entry(name.clone()).or_default();
+                for key in [kind.to_string(), "all".to_string()] {
+                    let t = table.entry(key).or_default();
+                    match fix {
+                        None => t.left += 1,
+                        Some(f) if f.eq_ignore_ascii_case(truth) => t.right += 1,
+                        Some(_) => t.wrong += 1,
+                    }
+                }
+            }
+            if let Some(f) = &pick_fix {
+                if !f.eq_ignore_ascii_case(truth) {
+                    let ctx: String = before.chars().rev().take(40).collect::<Vec<_>>().into_iter().rev().collect();
+                    wrong_lines.push(format!(
+                        "  {kind:9} …{ctx} [{typo}] -> {f:?} (meant {truth:?}; 0.1.3: {base_fix:?})"
+                    ));
+                }
+            }
+        }
+    }
+    println!("{samples} typos");
+    for (name, table) in &tables {
+        let all = &table["all"];
+        println!("{name:36} right {:4}  wrong {:4}  left {:4}", all.right, all.wrong, all.left);
+    }
+    for name in ["0.1.3 (rules + veto)", "model picks"] {
+        println!("{name}:");
+        for (kind, t) in &tables[name] {
+            println!("  {kind:9} right {:4}  wrong {:4}  left {:4}", t.right, t.wrong, t.left);
+        }
+    }
+    println!("wrong fixes with model picks:");
+    for line in wrong_lines {
+        println!("{line}");
+    }
+    Ok(())
+}

crates/cc-app/src/fixgate.rs

@@ -149,6 +149,77 @@ pub fn model_prefers(typo: &str, fix: &str, top: &[(String, f32)], is_word: impl
     (lp - fix_lp >= 4f32.ln()).then_some(word)
 }
 
+/// How likely the word [`model_pick`] returns must be as the next token, at least.
+const PICK_FLOOR: f32 = -4.6; // ln 0.01
+/// Nats of model likelihood that one unit of typing-error cost is worth in [`model_pick`].
+const COST_WEIGHT: f32 = 2.0;
+
+/// The word the language model clearly expects where `typo` was typed, among the words the
+/// typo may stand for (`options`: word and typing-error cost, from `Checker::typo_options`),
+/// given the model's likeliest next tokens (`top`) after the text before the typo. Each
+/// option scores how likely the model finds it next, as a whole word token (` month`), less
+/// what its typo costs, so a likely word the keys do not explain loses. The best must be
+/// in the list, and score four times higher than anything else that could be meant there:
+/// another option, a token that starts the typo or an option (` doc` may begin what was
+/// typed), or the word as typed, which is as likely as the last token in the list for all
+/// the model says. `this morth` gives `month`. After `for the`, `calory` gives nothing,
+/// since the model expects neither `calorie` nor `color` there.
+pub fn model_pick(typo: &str, options: &[(String, f32)], top: &[(String, f32)]) -> Option<String> {
+    model_pick_tuned(typo, options, top, COST_WEIGHT, true)
+}
+
+/// [`model_pick`] with its cost weight, and with or without the word as typed as a rival
+/// (`typed_rival`), for tuning (examples/fix_eval.rs).
+pub fn model_pick_tuned(
+    typo: &str,
+    options: &[(String, f32)],
+    top: &[(String, f32)],
+    weight: f32,
+    typed_rival: bool,
+) -> Option<String> {
+    let typo = typo.to_lowercase();
+    let floor = top.iter().map(|(_, lp)| *lp).reduce(f32::min)?;
+    // Whole-word tokens; one without a leading space would continue the word before the typo.
+    let words: Vec<(String, f32)> = top
+        .iter()
+        .filter_map(|(token, lp)| {
+            let word = token.to_lowercase().strip_prefix(' ')?.to_string();
+            (!word.is_empty() && word.chars().all(char::is_alphabetic)).then_some((word, *lp))
+        })
+        .collect();
+    let listed = |w: &str| words.iter().find(|(t, _)| t == w).map(|(_, lp)| *lp);
+    let options: Vec<(String, f32)> = options.iter().map(|(o, cost)| (o.to_lowercase(), *cost)).collect();
+    // (option, its token's log probability if listed, score); an option not in the list is
+    // at most as likely as the floor.
+    let scored: Vec<(&str, Option<f32>, f32)> = options
+        .iter()
+        .map(|(o, cost)| {
+            let lp = listed(o);
+            (o.as_str(), lp, lp.unwrap_or(floor) - weight * cost)
+        })
+        .collect();
+    let (best_at, &(word, lp, score)) = scored
+        .iter()
+        .enumerate()
+        .filter(|(_, (_, lp, _))| lp.is_some())
+        .max_by(|a, b| a.1 .2.total_cmp(&b.1 .2))?;
+    let mut rival = if typed_rival { floor } else { f32::NEG_INFINITY };
+    for (at, (_, _, other)) in scored.iter().enumerate() {
+        if at != best_at {
+            rival = rival.max(*other);
+        }
+    }
+    for (token, token_lp) in &words {
+        let is_option = options.iter().any(|(o, _)| o == token);
+        let starts = typo.starts_with(token.as_str()) || options.iter().any(|(o, _)| o.starts_with(token.as_str()));
+        if !is_option && token.chars().count() >= 2 && starts {
+            rival = rival.max(*token_lp);
+        }
+    }
+    let lp = lp?;
+    (lp >= PICK_FLOOR && score - rival >= 4f32.ln()).then(|| word.to_string())
+}
+
 /// The end of `s` with letters written as `a` or `A` and digits as `0`, for the log: it
 /// shows a joined or doubled word without showing what was typed.
 pub fn shape(s: &str) -> String {
@@ -449,6 +520,31 @@ mod tests {
         );
     }
 
+    #[test]
+    fn model_pick_takes_the_option_the_model_clearly_expects() {
+        let top = |pairs: &[(&str, f32)]| pairs.iter().map(|(t, lp)| (t.to_string(), *lp)).collect::<Vec<_>>();
+        let options = |words: &[(&str, f32)]| words.iter().map(|(w, c)| (w.to_string(), *c)).collect::<Vec<_>>();
+        let morth = options(&[("north", 0.3), ("moth", 0.6), ("month", 0.9), ("worth", 1.4)]);
+        let after_this = top(&[(" month", -1.2), (" week", -1.9), (" year", -2.0), (" north", -5.0), (" time", -6.0)]);
+        assert_eq!(model_pick("morth", &morth, &after_this).as_deref(), Some("month"));
+        // Neither option is expected: nothing below the list can be ruled out.
+        let after_the = top(&[(" meeting", -1.4), (" next", -2.7), (" first", -2.8), (" time", -3.9)]);
+        assert_eq!(model_pick("calory", &options(&[("calorie", 1.0)]), &after_the), None);
+        // Close to another option once the typing cost counts.
+        let close = top(&[(" month", -1.2), (" moth", -1.9), (" time", -6.0)]);
+        assert_eq!(model_pick("morth", &morth, &close), None);
+        // Likely, but the keys do not explain it well enough.
+        let costly = top(&[(" worth", -1.5), (" time", -4.5)]);
+        assert_eq!(model_pick("morth", &morth, &costly), None);
+        // A token the typo starts with may be the start of the word as typed.
+        let typed = top(&[(" document", -2.0), (" doc", -2.5), (" time", -6.0)]);
+        assert_eq!(model_pick("documnet", &options(&[("document", 0.4)]), &typed), None);
+        // Too unlikely on its own, and a token that continues the word before is not a word.
+        assert_eq!(model_pick("morth", &morth, &top(&[(" month", -5.0), (" a", -9.0)])), None);
+        assert_eq!(model_pick("morth", &morth, &top(&[("month", -0.5), (" a", -6.0)])), None);
+        assert_eq!(model_pick("morth", &morth, &[]), None);
+    }
+
     #[test]
     fn shape_hides_letters_and_digits_only() {
         assert_eq!(shape("Previous the 2nd, ok."), "Aaaaaaaa aaa 0aa, aa.");