mirror of
https://github.com/openai/codex.git
synced 2026-11-13 14:46:14 +00:00
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:
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,
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in new issue
Block a user