mirror of
https://github.com/openai/codex.git
synced 2026-10-18 17:18:52 +00:00
Add multi-query lexical skill selection (#33614)
## What changed - Add a lexical selector that splits compound queries into sentence and connector-delimited views. - Merge candidates across the full query and its component views, prioritizing their best rank while using full-query rank, view coverage, and stable identifiers as tie-breakers. - Bound query decomposition and result counts, propagate truncation signals, and include the selector in the shadow-selection experiment. ## Testing - Add unit coverage for compound-query ranking, single-query compatibility, query splitting, and bounded-input signals. GitOrigin-RevId: 2442bf7198fc24295cbc270b9d222275ae8a7cc3
This commit is contained in:
@@ -1,8 +1,10 @@
|
||||
mod character_ngram;
|
||||
mod fielded_bm25;
|
||||
mod multi_query_lexical;
|
||||
mod weighted_lexical;
|
||||
pub(crate) use character_ngram::CharacterNgramSkillSelector;
|
||||
pub(crate) use fielded_bm25::FieldedBm25SkillSelector;
|
||||
pub(crate) use multi_query_lexical::MultiQueryLexicalSkillSelector;
|
||||
pub(crate) use weighted_lexical::WeightedLexicalSkillSelector;
|
||||
|
||||
/// Metadata searched by a cheap skill selector.
|
||||
|
||||
@@ -0,0 +1,166 @@
|
||||
use std::collections::HashMap;
|
||||
use std::collections::HashSet;
|
||||
|
||||
use super::CheapSkillSelection;
|
||||
use super::CheapSkillSelector;
|
||||
use super::SkillSelectionDocument;
|
||||
use super::WeightedLexicalSkillSelector;
|
||||
|
||||
const MAX_QUERY_VIEWS: usize = 8;
|
||||
const MAX_DECOMPOSITION_BYTES: usize = 4 * 1024;
|
||||
const MAX_RESULTS: usize = 50;
|
||||
const CONNECTORS: &[&str] = &[" and then ", " and ", " then ", " also "];
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default)]
|
||||
pub(crate) struct MultiQueryLexicalSkillSelector;
|
||||
|
||||
impl CheapSkillSelector for MultiQueryLexicalSkillSelector {
|
||||
fn method(&self) -> &'static str {
|
||||
"multi_query_lexical_v1"
|
||||
}
|
||||
|
||||
fn select(
|
||||
&self,
|
||||
query: &str,
|
||||
documents: &[SkillSelectionDocument<'_>],
|
||||
limit: usize,
|
||||
) -> CheapSkillSelection {
|
||||
let full_selection =
|
||||
WeightedLexicalSkillSelector.select(query, documents, limit.min(MAX_RESULTS));
|
||||
let views = query_views(query);
|
||||
if views.len() <= 1 || limit == 0 {
|
||||
return full_selection;
|
||||
}
|
||||
|
||||
let mut candidates = HashMap::new();
|
||||
record_candidates(&mut candidates, &full_selection, /*view_index*/ 0);
|
||||
let mut query_truncated = full_selection.query_truncated;
|
||||
let mut candidate_set_truncated = full_selection.candidate_set_truncated;
|
||||
for (view_index, view) in views.into_iter().enumerate().skip(1) {
|
||||
let selection =
|
||||
WeightedLexicalSkillSelector.select(view, documents, limit.min(MAX_RESULTS));
|
||||
query_truncated |= selection.query_truncated;
|
||||
candidate_set_truncated |= selection.candidate_set_truncated;
|
||||
record_candidates(&mut candidates, &selection, view_index);
|
||||
}
|
||||
|
||||
let mut candidates = candidates.into_values().collect::<Vec<_>>();
|
||||
candidates.sort_by(|left, right| {
|
||||
left.best_rank
|
||||
.cmp(&right.best_rank)
|
||||
.then_with(|| {
|
||||
left.full_query_rank
|
||||
.unwrap_or(usize::MAX)
|
||||
.cmp(&right.full_query_rank.unwrap_or(usize::MAX))
|
||||
})
|
||||
.then_with(|| right.view_count.cmp(&left.view_count))
|
||||
.then_with(|| left.first_view.cmp(&right.first_view))
|
||||
.then_with(|| left.id.cmp(&right.id))
|
||||
});
|
||||
|
||||
CheapSkillSelection {
|
||||
candidate_ids: candidates
|
||||
.into_iter()
|
||||
.take(limit.min(MAX_RESULTS))
|
||||
.map(|candidate| candidate.id)
|
||||
.collect(),
|
||||
query_term_count: full_selection.query_term_count,
|
||||
query_truncated,
|
||||
candidate_set_truncated,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct RankedCandidate {
|
||||
id: usize,
|
||||
best_rank: usize,
|
||||
full_query_rank: Option<usize>,
|
||||
view_count: usize,
|
||||
first_view: usize,
|
||||
}
|
||||
|
||||
fn record_candidates(
|
||||
candidates: &mut HashMap<usize, RankedCandidate>,
|
||||
selection: &CheapSkillSelection,
|
||||
view_index: usize,
|
||||
) {
|
||||
for (rank, id) in selection.candidate_ids.iter().copied().enumerate() {
|
||||
let rank = rank + 1;
|
||||
candidates
|
||||
.entry(id)
|
||||
.and_modify(|candidate| {
|
||||
candidate.best_rank = candidate.best_rank.min(rank);
|
||||
candidate.view_count = candidate.view_count.saturating_add(1);
|
||||
if view_index == 0 {
|
||||
candidate.full_query_rank = Some(rank);
|
||||
}
|
||||
})
|
||||
.or_insert(RankedCandidate {
|
||||
id,
|
||||
best_rank: rank,
|
||||
full_query_rank: (view_index == 0).then_some(rank),
|
||||
view_count: 1,
|
||||
first_view: view_index,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
fn query_views(query: &str) -> Vec<&str> {
|
||||
let full_query = bounded(query, MAX_DECOMPOSITION_BYTES).trim();
|
||||
if full_query.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let mut views = vec![full_query];
|
||||
let mut seen = HashSet::from([full_query]);
|
||||
for sentence in full_query.split(['\n', '\r', '.', '!', '?', ';']) {
|
||||
for clause in split_connectors(sentence) {
|
||||
let clause = clause.trim();
|
||||
if clause.chars().count() >= 2 && seen.insert(clause) {
|
||||
views.push(clause);
|
||||
if views.len() == MAX_QUERY_VIEWS {
|
||||
return views;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
views
|
||||
}
|
||||
|
||||
fn bounded(value: &str, max_bytes: usize) -> &str {
|
||||
if value.len() <= max_bytes {
|
||||
return value;
|
||||
}
|
||||
let mut end = max_bytes;
|
||||
while !value.is_char_boundary(end) {
|
||||
end = end.saturating_sub(1);
|
||||
}
|
||||
&value[..end]
|
||||
}
|
||||
|
||||
fn split_connectors(value: &str) -> Vec<&str> {
|
||||
let lowercase = value.to_ascii_lowercase();
|
||||
let mut segments = Vec::new();
|
||||
let mut start = 0;
|
||||
while start < value.len() {
|
||||
let next = CONNECTORS
|
||||
.iter()
|
||||
.filter_map(|connector| {
|
||||
lowercase[start..]
|
||||
.find(connector)
|
||||
.map(|offset| (start + offset, connector.len()))
|
||||
})
|
||||
.min_by_key(|(position, _)| *position);
|
||||
let Some((position, connector_length)) = next else {
|
||||
break;
|
||||
};
|
||||
segments.push(&value[start..position]);
|
||||
start = position + connector_length;
|
||||
}
|
||||
segments.push(&value[start..]);
|
||||
segments
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "multi_query_lexical_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,90 @@
|
||||
use super::*;
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
#[test]
|
||||
fn multi_query_promotes_each_clause_leader() {
|
||||
let documents = [
|
||||
document(/*id*/ 1, "rust-format", "Format Rust source code."),
|
||||
document(/*id*/ 2, "rust-lint", "Fix Rust source lint errors."),
|
||||
document(/*id*/ 3, "rust-review", "Review Rust source code."),
|
||||
document(
|
||||
/*id*/ 4,
|
||||
"ci-fix",
|
||||
"Diagnose failing GitHub Actions checks.",
|
||||
),
|
||||
];
|
||||
|
||||
let selection = MultiQueryLexicalSkillSelector.select(
|
||||
"format and review Rust source code, and then diagnose failing GitHub Actions checks",
|
||||
&documents,
|
||||
/*limit*/ 4,
|
||||
);
|
||||
|
||||
assert!(selection.candidate_ids[..3].contains(&1));
|
||||
assert!(selection.candidate_ids[..3].contains(&4));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn single_query_matches_the_underlying_selector() {
|
||||
let documents = [
|
||||
document(/*id*/ 1, "presentations", "Create visual decks."),
|
||||
document(/*id*/ 2, "spreadsheets", "Analyze tabular data."),
|
||||
];
|
||||
|
||||
let expected =
|
||||
WeightedLexicalSkillSelector.select("create presentations", &documents, /*limit*/ 20);
|
||||
let actual = MultiQueryLexicalSkillSelector.select(
|
||||
"create presentations",
|
||||
&documents,
|
||||
/*limit*/ 20,
|
||||
);
|
||||
|
||||
assert_eq!(expected, actual);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn query_views_split_sentences_and_connectors() {
|
||||
assert_eq!(
|
||||
vec![
|
||||
"format code and then fix tests; write a summary",
|
||||
"format code",
|
||||
"fix tests",
|
||||
"write a summary",
|
||||
],
|
||||
query_views("format code and then fix tests; write a summary"),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multi_query_preserves_bounded_input_signals() {
|
||||
let long_query = format!("{}\nand inspect logs", "match ".repeat(4 * 1024));
|
||||
let names = (0..=1_000)
|
||||
.map(|index| format!("candidate-{index}"))
|
||||
.collect::<Vec<_>>();
|
||||
let documents = names
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(id, name)| SkillSelectionDocument {
|
||||
id,
|
||||
name,
|
||||
short_description: None,
|
||||
description: "match logs",
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let selection =
|
||||
MultiQueryLexicalSkillSelector.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,
|
||||
}
|
||||
}
|
||||
@@ -15,6 +15,7 @@ use crate::dynamic_skill_selector::CharacterNgramSkillSelector;
|
||||
use crate::dynamic_skill_selector::CheapSkillSelection;
|
||||
use crate::dynamic_skill_selector::CheapSkillSelector;
|
||||
use crate::dynamic_skill_selector::FieldedBm25SkillSelector;
|
||||
use crate::dynamic_skill_selector::MultiQueryLexicalSkillSelector;
|
||||
use crate::dynamic_skill_selector::SkillSelectionDocument;
|
||||
use crate::dynamic_skill_selector::WeightedLexicalSkillSelector;
|
||||
|
||||
@@ -41,6 +42,7 @@ impl ShadowSelectionExperiment {
|
||||
Box::new(WeightedLexicalSkillSelector),
|
||||
Box::new(FieldedBm25SkillSelector),
|
||||
Box::new(CharacterNgramSkillSelector),
|
||||
Box::new(MultiQueryLexicalSkillSelector),
|
||||
],
|
||||
metrics_client,
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user