///|
pub fn deduplicate_candidate_pool(pool : CandidatePool) -> CandidatePool {
let seen : Map[String, Unit] = Map([])
let docs : Array[String] = []
for doc_id in pool.doc_ids {
if !seen.contains(doc_id) {
seen[doc_id] = ()
docs.push(doc_id)
}
}
{ query_id: pool.query_id, doc_ids: docs }
}
///|
pub fn interleave_candidate_pools(
left : CandidatePool,
right : CandidatePool,
) -> CandidatePool {
let left_pool = deduplicate_candidate_pool(left)
let right_pool = deduplicate_candidate_pool(right)
let result : Array[String] = []
let seen : Map[String, Unit] = Map([])
let limit = Int::max(left_pool.doc_ids.length(), right_pool.doc_ids.length())
for index in 0.. Map[String, Unit] {
let blocked : Map[String, Unit] = Map([])
for item in qrels {
if item.query_id == query_id && (skip_judged || item.relevance >= threshold) {
blocked[item.doc_id] = ()
}
}
blocked
}
///|
pub fn sample_pool_negatives(
qrels : Array[JudgedDoc],
pools : Array[CandidatePool],
config? : NegativeSampleConfig = NegativeSampleConfig::default(),
) -> Array[NegativeSample] {
let result : Array[NegativeSample] = []
for pool in pools {
let clean_pool = deduplicate_candidate_pool(pool)
let blocked = judged_doc_set(
qrels,
clean_pool.query_id,
config.relevant_threshold,
config.skip_judged,
)
let selected = sampled_doc_ids(
clean_pool.doc_ids,
config.strategy,
Int::max(config.per_query, 0) * 4,
)
let seen : Map[String, Unit] = Map([])
for doc_rank in selected {
let (doc_id, source_rank) = doc_rank
if blocked.contains(doc_id) || seen.contains(doc_id) {
continue
}
seen[doc_id] = ()
result.push({
query_id: clean_pool.query_id,
doc_id,
source_rank,
strategy: strategy_name(config.strategy),
})
if seen.length() >= Int::max(config.per_query, 0) {
break
}
}
}
result.sort_by(fn(a, b) {
let query_order = a.query_id.compare(b.query_id)
if query_order == 0 {
a.source_rank - b.source_rank
} else {
query_order
}
})
result
}
///|
pub fn sample_pool_negatives_tsv(
qrels : Array[JudgedDoc],
pools : Array[CandidatePool],
config? : NegativeSampleConfig = NegativeSampleConfig::default(),
) -> String {
render_negative_samples_tsv(sample_pool_negatives(qrels, pools, config~))
}
///|
pub fn sampling_strategy_description(strategy : NegativeStrategy) -> String {
match strategy {
Tail(window) =>
"tail: choose from the last \{Int::max(window, 1)} candidate ranks"
HardWindow(window) =>
"hard: choose from the first \{Int::max(window, 1)} candidate ranks"
Stride(step) => "stride: choose every \{Int::max(step, 1)} candidate rank"
}
}
///|
pub fn sampling_plan(
qrels : Array[JudgedDoc],
pools : Array[CandidatePool],
config? : NegativeSampleConfig = NegativeSampleConfig::default(),
) -> String {
let profile = profile_candidate_pools(pools)
let coverage = candidate_pool_coverage(qrels, pools)
let lines : Array[String] = [
"strategy=\{sampling_strategy_description(config.strategy)}",
"per_query=\{Int::max(config.per_query, 0)}",
"skip_judged=\{config.skip_judged}",
"pools=\{profile.pool_count}",
"candidate_coverage=\{format_metric(coverage)}",
]
lines.join("\n")
}