Add character n-gram skill selection (#33613)

## What changed

- Add a character n-gram selector that scores matches across skill names, short descriptions, and descriptions using field weights and inverse document frequency.
- Include the selector in the shadow selection experiment alongside the existing lexical selectors.
- Bound query, document, candidate, and result processing, and report query or candidate truncation through the existing selection metadata.

## Testing

Add unit coverage for related word forms, typos, CJK text without word boundaries, and bounded inputs.

GitOrigin-RevId: 4df60cc661338050eeb7741c716bc1c9f730fedd
This commit is contained in:
jif authored and copyberry committed 2026-07-16 14:38:36 +00:00
1 parent a47c661ea9
commit c983a53f20
4 files changed
+324

No files matched your search

@@ -1,5 +1,7 @@
mod character_ngram;
mod fielded_bm25; mod fielded_bm25;
mod weighted_lexical; mod weighted_lexical;
pub(crate) use character_ngram::CharacterNgramSkillSelector;
pub(crate) use fielded_bm25::FieldedBm25SkillSelector; pub(crate) use fielded_bm25::FieldedBm25SkillSelector;
pub(crate) use weighted_lexical::WeightedLexicalSkillSelector; pub(crate) use weighted_lexical::WeightedLexicalSkillSelector;
@@ -0,0 +1,243 @@
use std::collections::HashMap;
use std::collections::HashSet;
use super::CheapSkillSelection;
use super::CheapSkillSelector;
use super::SkillSelectionDocument;
const MAX_QUERY_BYTES: usize = 4 * 1024;
const MAX_QUERY_TERMS: usize = 64;
const MAX_QUERY_GRAMS: usize = 512;
const MAX_DOCUMENT_BYTES: usize = 4 * 1024;
const MAX_DOCUMENT_TERMS: usize = 256;
const MAX_DOCUMENT_GRAMS: usize = 512;
const MAX_CANDIDATES: usize = 1_000;
const MAX_RESULTS: usize = 50;
const MIN_GRAM_CHARS: usize = 2;
const MAX_GRAM_CHARS: usize = 5;
const FIELD_WEIGHTS: [f64; 3] = [8.0, 4.0, 1.0];
const STOP_WORDS: &[&str] = &[
"a", "an", "and", "are", "as", "at", "be", "by", "do", "for", "from", "how", "i", "in", "is",
"it", "me", "my", "of", "on", "or", "please", "that", "the", "this", "to", "use", "we", "what",
"when", "where", "which", "with", "you", "your",
];
#[derive(Clone, Copy, Debug, Default)]
pub(crate) struct CharacterNgramSkillSelector;
impl CheapSkillSelector for CharacterNgramSkillSelector {
fn method(&self) -> &'static str {
"character_ngram_v1"
}
fn select(
&self,
query: &str,
documents: &[SkillSelectionDocument<'_>],
limit: usize,
) -> CheapSkillSelection {
let (query, query_bytes_truncated) = bounded(query, MAX_QUERY_BYTES);
let (query_terms, query_terms_truncated) = query_terms(query);
let (query_grams, query_grams_truncated) = grams(&query_terms, MAX_QUERY_GRAMS);
let query_truncated =
query_bytes_truncated || query_terms_truncated || query_grams_truncated;
let candidate_set_truncated = documents.len() > MAX_CANDIDATES;
if query_grams.is_empty() || limit == 0 {
return CheapSkillSelection {
query_term_count: query_terms.len(),
query_truncated,
candidate_set_truncated,
..Default::default()
};
}
let prepared = documents
.iter()
.take(MAX_CANDIDATES)
.map(PreparedDocument::new)
.collect::<Vec<_>>();
let document_frequencies = document_frequencies(&prepared);
let document_count = prepared.len() as f64;
let minimum_matches = query_grams.len().min(3);
let mut scored = prepared
.iter()
.filter_map(|document| {
let (score, matched_grams) = score_document(
document,
&query_grams,
&document_frequencies,
document_count,
);
(matched_grams >= minimum_matches).then_some((score, document.id, document.name))
})
.collect::<Vec<_>>();
scored.sort_by(|left, right| {
right
.0
.total_cmp(&left.0)
.then_with(|| left.2.cmp(right.2))
.then_with(|| left.1.cmp(&right.1))
});
CheapSkillSelection {
candidate_ids: scored
.into_iter()
.take(limit.min(MAX_RESULTS))
.map(|(_, id, _)| id)
.collect(),
query_term_count: query_terms.len(),
query_truncated,
candidate_set_truncated,
}
}
}
struct PreparedDocument<'a> {
id: usize,
name: &'a str,
fields: [HashSet<String>; 3],
}
impl<'a> PreparedDocument<'a> {
fn new(document: &'a SkillSelectionDocument<'a>) -> Self {
Self {
id: document.id,
name: document.name,
fields: [
document_grams(document.name),
document_grams(document.short_description.unwrap_or_default()),
document_grams(document.description),
],
}
}
}
fn score_document(
document: &PreparedDocument<'_>,
query_grams: &[String],
document_frequencies: &HashMap<String, usize>,
document_count: f64,
) -> (f64, usize) {
query_grams.iter().fold((0.0, 0), |(score, matched), gram| {
let frequency = document_frequencies.get(gram).copied().unwrap_or_default() as f64;
if frequency == 0.0 {
return (score, matched);
}
let inverse_document_frequency =
(1.0 + (document_count - frequency + 0.5) / (frequency + 0.5)).ln();
let field_weight = document
.fields
.iter()
.enumerate()
.filter(|(_, field)| field.contains(gram))
.map(|(index, _)| FIELD_WEIGHTS[index])
.sum::<f64>();
if field_weight == 0.0 {
(score, matched)
} else {
(
score + inverse_document_frequency * field_weight,
matched + 1,
)
}
})
}
fn document_frequencies(documents: &[PreparedDocument<'_>]) -> HashMap<String, usize> {
let mut frequencies = HashMap::new();
for document in documents {
let grams = document
.fields
.iter()
.flatten()
.map(String::as_str)
.collect::<HashSet<_>>();
for gram in grams {
*frequencies.entry(gram.to_string()).or_default() += 1;
}
}
frequencies
}
fn query_terms(query: &str) -> (Vec<String>, bool) {
let mut seen = HashSet::new();
let mut terms = Vec::new();
for term in normalized_terms(query)
.into_iter()
.filter(|term| term.chars().count() >= 2 && !STOP_WORDS.contains(&term.as_str()))
{
if !seen.insert(term.clone()) {
continue;
}
if terms.len() == MAX_QUERY_TERMS {
return (terms, true);
}
terms.push(term);
}
(terms, false)
}
fn document_grams(value: &str) -> HashSet<String> {
let (value, _) = bounded(value, MAX_DOCUMENT_BYTES);
let terms = normalized_terms(value)
.into_iter()
.take(MAX_DOCUMENT_TERMS)
.collect::<Vec<_>>();
grams(&terms, MAX_DOCUMENT_GRAMS).0.into_iter().collect()
}
fn grams(terms: &[String], limit: usize) -> (Vec<String>, bool) {
let mut seen = HashSet::new();
let mut grams = Vec::new();
for term in terms {
let characters = term.chars().collect::<Vec<_>>();
let minimum_gram_chars = if term.is_ascii() && characters.len() > MIN_GRAM_CHARS {
MIN_GRAM_CHARS + 1
} else {
MIN_GRAM_CHARS
};
for gram_size in minimum_gram_chars..=MAX_GRAM_CHARS.min(characters.len()) {
for start in 0..=characters.len() - gram_size {
let gram = characters[start..start + gram_size]
.iter()
.collect::<String>();
if !seen.insert(gram.clone()) {
continue;
}
if grams.len() == limit {
return (grams, true);
}
grams.push(gram);
}
}
}
(grams, false)
}
fn normalized_terms(value: &str) -> Vec<String> {
let mut normalized = String::with_capacity(value.len());
for character in value.chars() {
if character.is_alphanumeric() {
normalized.extend(character.to_lowercase());
} else {
normalized.push(' ');
}
}
normalized.split_whitespace().map(str::to_string).collect()
}
fn bounded(value: &str, max_bytes: usize) -> (&str, bool) {
if value.len() <= max_bytes {
return (value, false);
}
let mut end = max_bytes;
while !value.is_char_boundary(end) {
end = end.saturating_sub(1);
}
(&value[..end], true)
}
#[cfg(test)]
#[path = "character_ngram_tests.rs"]
mod tests;
@@ -0,0 +1,77 @@
use super::*;
use pretty_assertions::assert_eq;
#[test]
fn ngrams_match_related_word_forms() {
let documents = [
document(/*id*/ 1, "presentations", "Create visual decks."),
document(/*id*/ 2, "spreadsheets", "Analyze tabular data."),
];
let selection =
CharacterNgramSkillSelector.select("create a presentation", &documents, /*limit*/ 20);
assert_eq!(vec![1], selection.candidate_ids);
}
#[test]
fn ngrams_tolerate_a_typo() {
let documents = [
document(/*id*/ 1, "postgresql", "Manage a relational database."),
document(/*id*/ 2, "postscript", "Render printable documents."),
];
let selection = CharacterNgramSkillSelector.select(
"repair my postgrez database",
&documents,
/*limit*/ 20,
);
assert_eq!(vec![1, 2], selection.candidate_ids);
}
#[test]
fn ngrams_match_cjk_without_word_boundaries() {
let documents = [
document(/*id*/ 1, "演示文稿", "创建幻灯片。"),
document(/*id*/ 2, "电子表格", "分析表格数据。"),
];
let selection =
CharacterNgramSkillSelector.select("帮我制作演示文稿", &documents, /*limit*/ 20);
assert_eq!(vec![1], selection.candidate_ids);
}
#[test]
fn ngrams_report_bounded_inputs() {
let long_query = "match ".repeat(MAX_QUERY_BYTES);
let names = (0..=MAX_CANDIDATES)
.map(|index| format!("candidate-{index}"))
.collect::<Vec<_>>();
let documents = names
.iter()
.enumerate()
.map(|(id, name)| SkillSelectionDocument {
id,
name,
short_description: None,
description: "match",
})
.collect::<Vec<_>>();
let selection = CharacterNgramSkillSelector.select(&long_query, &documents, /*limit*/ 20);
assert!(selection.query_truncated);
assert!(selection.candidate_set_truncated);
assert_eq!(20, selection.candidate_ids.len());
}
fn document<'a>(id: usize, name: &'a str, description: &'a str) -> SkillSelectionDocument<'a> {
SkillSelectionDocument {
id,
name,
short_description: None,
description,
}
}
@@ -11,6 +11,7 @@ use codex_protocol::user_input::UserInput;
use crate::catalog::SkillCatalog; use crate::catalog::SkillCatalog;
use crate::catalog::SkillSourceKind; use crate::catalog::SkillSourceKind;
use crate::dynamic_skill_selector::CharacterNgramSkillSelector;
use crate::dynamic_skill_selector::CheapSkillSelection; use crate::dynamic_skill_selector::CheapSkillSelection;
use crate::dynamic_skill_selector::CheapSkillSelector; use crate::dynamic_skill_selector::CheapSkillSelector;
use crate::dynamic_skill_selector::FieldedBm25SkillSelector; use crate::dynamic_skill_selector::FieldedBm25SkillSelector;
@@ -39,6 +40,7 @@ impl ShadowSelectionExperiment {
selectors: vec![ selectors: vec![
Box::new(WeightedLexicalSkillSelector), Box::new(WeightedLexicalSkillSelector),
Box::new(FieldedBm25SkillSelector), Box::new(FieldedBm25SkillSelector),
Box::new(CharacterNgramSkillSelector),
], ],
metrics_client, metrics_client,
} }