mirror of
https://github.com/openai/codex.git
synced 2026-10-18 17:18:52 +00:00
Add reciprocal rank fusion skill selection (#34547)
## What changed - Add a skill selector that uses reciprocal rank fusion to combine the top 50 results from the weighted lexical and character n-gram selectors. - Prefer candidates supported by both rankings, with deterministic rank and identifier tie-breakers, and propagate truncation metadata from both selectors. - Include the fused selector in the shadow-selection experiment under `rrf_lexical_char_v1`. ## Testing - Add unit coverage for cross-ranking support, deterministic tie-breaking, and combined input-bound reporting. GitOrigin-RevId: 3165136b8729ea3a9484f25c7f2ce0f91fb8aaa9
This commit is contained in:
@@ -1,10 +1,12 @@
|
||||
mod character_ngram;
|
||||
mod fielded_bm25;
|
||||
mod multi_query_lexical;
|
||||
mod rrf_lexical_char;
|
||||
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 rrf_lexical_char::RrfLexicalCharSkillSelector;
|
||||
pub(crate) use weighted_lexical::WeightedLexicalSkillSelector;
|
||||
|
||||
/// Metadata searched by a cheap skill selector.
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
use super::CharacterNgramSkillSelector;
|
||||
use super::CheapSkillSelection;
|
||||
use super::CheapSkillSelector;
|
||||
use super::SkillSelectionDocument;
|
||||
use super::WeightedLexicalSkillSelector;
|
||||
|
||||
const MAX_FUSION_CANDIDATES: usize = 50;
|
||||
const RRF_K: usize = 60;
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default)]
|
||||
pub(crate) struct RrfLexicalCharSkillSelector;
|
||||
|
||||
impl CheapSkillSelector for RrfLexicalCharSkillSelector {
|
||||
fn method(&self) -> &'static str {
|
||||
"rrf_lexical_char_v1"
|
||||
}
|
||||
|
||||
fn select(
|
||||
&self,
|
||||
query: &str,
|
||||
documents: &[SkillSelectionDocument<'_>],
|
||||
limit: usize,
|
||||
) -> CheapSkillSelection {
|
||||
if limit == 0 {
|
||||
return CheapSkillSelection::default();
|
||||
}
|
||||
|
||||
let lexical = WeightedLexicalSkillSelector.select(
|
||||
query,
|
||||
documents,
|
||||
/*limit*/ MAX_FUSION_CANDIDATES,
|
||||
);
|
||||
let character = CharacterNgramSkillSelector.select(
|
||||
query,
|
||||
documents,
|
||||
/*limit*/ MAX_FUSION_CANDIDATES,
|
||||
);
|
||||
let candidate_ids = fuse_rankings(
|
||||
[&lexical.candidate_ids, &character.candidate_ids],
|
||||
limit.min(MAX_FUSION_CANDIDATES),
|
||||
);
|
||||
|
||||
CheapSkillSelection {
|
||||
candidate_ids,
|
||||
query_term_count: lexical.query_term_count.max(character.query_term_count),
|
||||
query_truncated: lexical.query_truncated || character.query_truncated,
|
||||
candidate_set_truncated: lexical.candidate_set_truncated
|
||||
|| character.candidate_set_truncated,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn fuse_rankings<const N: usize>(rankings: [&[usize]; N], limit: usize) -> Vec<usize> {
|
||||
let mut scores = HashMap::<usize, f64>::new();
|
||||
let mut best_ranks = HashMap::<usize, usize>::new();
|
||||
for ranking in rankings {
|
||||
for (index, id) in ranking.iter().copied().enumerate() {
|
||||
let rank = index + 1;
|
||||
*scores.entry(id).or_default() += 1.0 / (RRF_K + rank) as f64;
|
||||
best_ranks
|
||||
.entry(id)
|
||||
.and_modify(|best_rank| *best_rank = (*best_rank).min(rank))
|
||||
.or_insert(rank);
|
||||
}
|
||||
}
|
||||
|
||||
let mut candidates = scores.into_iter().collect::<Vec<_>>();
|
||||
candidates.sort_by(|(left_id, left_score), (right_id, right_score)| {
|
||||
right_score
|
||||
.total_cmp(left_score)
|
||||
.then_with(|| best_ranks[left_id].cmp(&best_ranks[right_id]))
|
||||
.then_with(|| left_id.cmp(right_id))
|
||||
});
|
||||
candidates
|
||||
.into_iter()
|
||||
.take(limit)
|
||||
.map(|(id, _)| id)
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "rrf_lexical_char_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,34 @@
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn fusion_prefers_candidates_supported_by_both_rankings() {
|
||||
let fused = fuse_rankings([&[1, 2, 3], &[2, 3, 4]], /*limit*/ 4);
|
||||
|
||||
assert_eq!(vec![2, 3, 1, 4], fused);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn fusion_uses_rank_and_id_to_break_ties() {
|
||||
let fused = fuse_rankings([&[20, 10], &[30]], /*limit*/ 3);
|
||||
|
||||
assert_eq!(vec![20, 30, 10], fused);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selector_reports_the_combined_input_bounds() {
|
||||
let documents = [SkillSelectionDocument {
|
||||
id: 7,
|
||||
name: "presentations",
|
||||
short_description: None,
|
||||
description: "Create visual decks.",
|
||||
}];
|
||||
let query = "presentation ".repeat(4 * 1024);
|
||||
|
||||
let selection = RrfLexicalCharSkillSelector.select(&query, &documents, /*limit*/ 20);
|
||||
|
||||
assert_eq!(vec![7], selection.candidate_ids);
|
||||
assert!(selection.query_truncated);
|
||||
assert!(!selection.candidate_set_truncated);
|
||||
}
|
||||
@@ -16,6 +16,7 @@ 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::RrfLexicalCharSkillSelector;
|
||||
use crate::dynamic_skill_selector::SkillSelectionDocument;
|
||||
use crate::dynamic_skill_selector::WeightedLexicalSkillSelector;
|
||||
|
||||
@@ -43,6 +44,7 @@ impl ShadowSelectionExperiment {
|
||||
Box::new(FieldedBm25SkillSelector),
|
||||
Box::new(CharacterNgramSkillSelector),
|
||||
Box::new(MultiQueryLexicalSkillSelector),
|
||||
Box::new(RrfLexicalCharSkillSelector),
|
||||
],
|
||||
metrics_client,
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user