#![forbid(unsafe_code)]
use kcode_speaker_dataset::{Dataset, OpenSetFold, active_rows};
use kcode_speaker_model::{Decision, FitInput, ModelConfig, ModelError, fit, identify};
use kcode_speaker_types::LabeledSample;
use std::cmp::Ordering;
use std::collections::{BTreeMap, BTreeSet};
use std::error::Error;
use std::fmt;
pub struct EvaluationPlan<'a> {
pub dataset: &'a Dataset,
pub tuning_folds: &'a [OpenSetFold],
pub held_out_fold: &'a OpenSetFold,
pub candidates: &'a [ModelConfig],
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct AggregateMetrics {
pub known_total: usize,
pub known_correct: usize,
pub known_misidentified: usize,
pub known_rejected: usize,
pub unknown_total: usize,
pub unknown_rejected: usize,
pub unknown_accepted: usize,
}
impl AggregateMetrics {
pub fn known_identification_rate(&self) -> f64 {
rate(self.known_correct, self.known_total)
}
pub fn known_misidentification_rate(&self) -> f64 {
rate(self.known_misidentified, self.known_total)
}
pub fn known_rejection_rate(&self) -> f64 {
rate(self.known_rejected, self.known_total)
}
pub fn unknown_rejection_rate(&self) -> f64 {
rate(self.unknown_rejected, self.unknown_total)
}
pub fn unknown_acceptance_rate(&self) -> f64 {
rate(self.unknown_accepted, self.unknown_total)
}
pub fn balanced_open_set_rate(&self) -> f64 {
(self.known_identification_rate() + self.unknown_rejection_rate()) / 2.0
}
fn merge(&mut self, other: Self) -> Result<(), EvalError> {
checked_add(&mut self.known_total, other.known_total)?;
checked_add(&mut self.known_correct, other.known_correct)?;
checked_add(&mut self.known_misidentified, other.known_misidentified)?;
checked_add(&mut self.known_rejected, other.known_rejected)?;
checked_add(&mut self.unknown_total, other.unknown_total)?;
checked_add(&mut self.unknown_rejected, other.unknown_rejected)?;
checked_add(&mut self.unknown_accepted, other.unknown_accepted)?;
Ok(())
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct CandidateEvaluation {
pub config: ModelConfig,
pub tuning: AggregateMetrics,
}
#[derive(Clone, Debug, PartialEq)]
pub struct EvaluationResult {
pub selected_config: ModelConfig,
pub candidate_evaluations: Vec<CandidateEvaluation>,
pub held_out: AggregateMetrics,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum EvalError {
EmptyCandidates,
EmptyTuningFolds,
DuplicateCandidate,
DuplicateTuningFold,
InvalidCandidate,
InvalidFold,
FoldOutsideOuterTraining,
GroupLeakage,
OpenSetViolation,
CountOverflow,
Model(ModelError),
}
impl fmt::Display for EvalError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
let text = match self {
Self::EmptyCandidates => "evaluation has no candidate configurations",
Self::EmptyTuningFolds => "evaluation has no tuning folds",
Self::DuplicateCandidate => "evaluation contains a duplicate candidate configuration",
Self::DuplicateTuningFold => "evaluation contains a duplicate tuning fold",
Self::InvalidCandidate => "candidate configuration is invalid",
Self::InvalidFold => "evaluation fold is malformed",
Self::FoldOutsideOuterTraining => {
"tuning fold contains a row outside held-out training"
}
Self::GroupLeakage => "a leakage group crosses training and test",
Self::OpenSetViolation => "fold violates known-speaker or pseudo-unknown boundaries",
Self::CountOverflow => "aggregate evaluation count overflowed",
Self::Model(error) => return error.fmt(formatter),
};
formatter.write_str(text)
}
}
impl Error for EvalError {
fn source(&self) -> Option<&(dyn Error + 'static)> {
match self {
Self::Model(error) => Some(error),
_ => None,
}
}
}
impl From<ModelError> for EvalError {
fn from(error: ModelError) -> Self {
Self::Model(error)
}
}
pub fn evaluate(plan: EvaluationPlan<'_>) -> Result<EvaluationResult, EvalError> {
if plan.candidates.is_empty() {
return Err(EvalError::EmptyCandidates);
}
if plan.tuning_folds.is_empty() {
return Err(EvalError::EmptyTuningFolds);
}
let mut candidates = plan.candidates.to_vec();
for candidate in &candidates {
validate_candidate_shape(candidate)?;
}
candidates.sort_by(compare_configs);
if candidates
.windows(2)
.any(|pair| compare_configs(&pair[0], &pair[1]).is_eq())
{
return Err(EvalError::DuplicateCandidate);
}
let rows = active_rows(plan.dataset);
let held_out = canonical_fold(plan.held_out_fold);
let complete_universe = vec![true; rows.len()];
validate_fold(rows, &held_out, &complete_universe, false)?;
let mut outer_training = vec![false; rows.len()];
for &index in &held_out.train_indices {
outer_training[index] = true;
}
let mut tuning_folds = plan
.tuning_folds
.iter()
.map(canonical_fold)
.collect::<Vec<_>>();
tuning_folds.sort_by(compare_folds);
if tuning_folds.windows(2).any(|pair| pair[0] == pair[1]) {
return Err(EvalError::DuplicateTuningFold);
}
for fold in &tuning_folds {
validate_fold(rows, fold, &outer_training, true)?;
}
let minimum_training_count = tuning_folds
.iter()
.map(|fold| fold.train_indices.len())
.chain(std::iter::once(held_out.train_indices.len()))
.min()
.ok_or(EvalError::EmptyTuningFolds)?;
if candidates
.iter()
.any(|candidate| usize::from(candidate.components) > minimum_training_count)
{
return Err(EvalError::InvalidCandidate);
}
let mut candidate_evaluations = Vec::with_capacity(candidates.len());
for config in candidates {
let mut tuning = AggregateMetrics::default();
for fold in &tuning_folds {
tuning.merge(score_fold(rows, fold, &config)?)?;
}
candidate_evaluations.push(CandidateEvaluation { config, tuning });
}
let mut selected_index = 0;
for index in 1..candidate_evaluations.len() {
if compare_evaluations(
&candidate_evaluations[index],
&candidate_evaluations[selected_index],
)
.is_lt()
{
selected_index = index;
}
}
let selected_config = candidate_evaluations[selected_index].config.clone();
let held_out_metrics = score_fold(rows, &held_out, &selected_config)?;
Ok(EvaluationResult {
selected_config,
candidate_evaluations,
held_out: held_out_metrics,
})
}
fn rate(numerator: usize, denominator: usize) -> f64 {
if denominator == 0 {
0.0
} else {
numerator as f64 / denominator as f64
}
}
fn checked_add(value: &mut usize, addition: usize) -> Result<(), EvalError> {
*value = value
.checked_add(addition)
.ok_or(EvalError::CountOverflow)?;
Ok(())
}
fn validate_candidate_shape(config: &ModelConfig) -> Result<(), EvalError> {
if config.components == 0
|| !config.relevance.is_finite()
|| config.relevance <= 0.0
|| !config.variance_floor.is_finite()
|| config.variance_floor <= 0.0
|| !config.absolute_threshold.is_finite()
|| !config.margin_threshold.is_finite()
{
return Err(EvalError::InvalidCandidate);
}
Ok(())
}
fn compare_configs(left: &ModelConfig, right: &ModelConfig) -> Ordering {
u64::from(left.mask)
.cmp(&u64::from(right.mask))
.then_with(|| left.components.cmp(&right.components))
.then_with(|| left.relevance.total_cmp(&right.relevance))
.then_with(|| left.variance_floor.total_cmp(&right.variance_floor))
.then_with(|| left.absolute_threshold.total_cmp(&right.absolute_threshold))
.then_with(|| left.margin_threshold.total_cmp(&right.margin_threshold))
}
fn canonical_fold(fold: &OpenSetFold) -> OpenSetFold {
let mut canonical = fold.clone();
canonical.train_indices.sort_unstable();
canonical.known_test_indices.sort_unstable();
canonical.unknown_test_indices.sort_unstable();
canonical
}
fn compare_folds(left: &OpenSetFold, right: &OpenSetFold) -> Ordering {
left.pseudo_unknown_speaker
.cmp(&right.pseudo_unknown_speaker)
.then_with(|| left.train_indices.cmp(&right.train_indices))
.then_with(|| left.known_test_indices.cmp(&right.known_test_indices))
.then_with(|| left.unknown_test_indices.cmp(&right.unknown_test_indices))
}
fn validate_fold(
rows: &[LabeledSample],
fold: &OpenSetFold,
universe: &[bool],
nested: bool,
) -> Result<(), EvalError> {
if universe.len() != rows.len()
|| fold.train_indices.is_empty()
|| fold.known_test_indices.is_empty()
|| fold.unknown_test_indices.is_empty()
{
return Err(EvalError::InvalidFold);
}
let mut placement = vec![0_u8; rows.len()];
for (side, indices) in [
(1_u8, fold.train_indices.as_slice()),
(2_u8, fold.known_test_indices.as_slice()),
(3_u8, fold.unknown_test_indices.as_slice()),
] {
for &index in indices {
if index >= rows.len() {
return Err(EvalError::InvalidFold);
}
if !universe[index] {
return Err(if nested {
EvalError::FoldOutsideOuterTraining
} else {
EvalError::InvalidFold
});
}
if placement[index] != 0 {
return Err(EvalError::InvalidFold);
}
placement[index] = side;
}
}
if universe
.iter()
.zip(&placement)
.any(|(included, side)| *included && *side == 0)
{
return Err(EvalError::InvalidFold);
}
let mut group_sides = BTreeMap::<&str, u8>::new();
for (index, row) in rows.iter().enumerate() {
if !universe[index] {
continue;
}
let side = if placement[index] == 1 { 1 } else { 2 };
match group_sides.get(row.group_id.as_ref()) {
Some(previous) if *previous != side => {
return Err(EvalError::GroupLeakage);
}
Some(_) => {}
None => {
group_sides.insert(row.group_id.as_ref(), side);
}
}
}
let pseudo_unknown = &fold.pseudo_unknown_speaker;
if fold
.train_indices
.iter()
.any(|&index| rows[index].speaker_id == *pseudo_unknown)
|| fold
.unknown_test_indices
.iter()
.any(|&index| rows[index].speaker_id != *pseudo_unknown)
|| rows.iter().enumerate().any(|(index, row)| {
universe[index] && row.speaker_id == *pseudo_unknown && placement[index] != 3
})
{
return Err(EvalError::OpenSetViolation);
}
let enrolled = fold
.train_indices
.iter()
.map(|&index| rows[index].speaker_id.as_ref())
.collect::<BTreeSet<_>>();
if fold
.known_test_indices
.iter()
.any(|&index| !enrolled.contains(rows[index].speaker_id.as_ref()))
{
return Err(EvalError::OpenSetViolation);
}
Ok(())
}
fn score_fold(
rows: &[LabeledSample],
fold: &OpenSetFold,
config: &ModelConfig,
) -> Result<AggregateMetrics, EvalError> {
let training = fold
.train_indices
.iter()
.map(|&index| rows[index].clone())
.collect::<Vec<_>>();
let cohort_id = training
.first()
.ok_or(EvalError::InvalidFold)?
.cohort_id
.clone();
let model = fit(FitInput {
cohort_id: &cohort_id,
samples: &training,
config: config.clone(),
})?;
let mut metrics = AggregateMetrics::default();
for &index in &fold.known_test_indices {
let row = &rows[index];
let result = identify(&model, &cohort_id, &row.features)?;
checked_add(&mut metrics.known_total, 1)?;
match result.decision {
Decision::Known { speaker_id } if speaker_id == row.speaker_id => {
checked_add(&mut metrics.known_correct, 1)?;
}
Decision::Known { .. } => {
checked_add(&mut metrics.known_misidentified, 1)?;
}
Decision::Unknown => {
checked_add(&mut metrics.known_rejected, 1)?;
}
}
}
for &index in &fold.unknown_test_indices {
let row = &rows[index];
let result = identify(&model, &cohort_id, &row.features)?;
checked_add(&mut metrics.unknown_total, 1)?;
match result.decision {
Decision::Unknown => {
checked_add(&mut metrics.unknown_rejected, 1)?;
}
Decision::Known { .. } => {
checked_add(&mut metrics.unknown_accepted, 1)?;
}
}
}
Ok(metrics)
}
fn compare_evaluations(left: &CandidateEvaluation, right: &CandidateEvaluation) -> Ordering {
right
.tuning
.balanced_open_set_rate()
.total_cmp(&left.tuning.balanced_open_set_rate())
.then_with(|| {
right
.tuning
.known_identification_rate()
.total_cmp(&left.tuning.known_identification_rate())
})
.then_with(|| {
right
.tuning
.unknown_rejection_rate()
.total_cmp(&left.tuning.unknown_rejection_rate())
})
.then_with(|| {
left.tuning
.known_misidentified
.cmp(&right.tuning.known_misidentified)
})
.then_with(|| compare_configs(&left.config, &right.config))
}
#[cfg(test)]
mod tests {
use super::*;
use kcode_speaker_dataset::build;
use kcode_speaker_model::ALL_24;
use kcode_speaker_types::{
FEATURE_COUNT, FeatureMask, FeatureVector, Key, ObjectId, RecordingKind,
};
fn key(value: &str) -> Key {
Key::parse(value).unwrap()
}
fn sample(
id: &str,
speaker: &str,
clip: &str,
group: &str,
first_feature: u8,
) -> LabeledSample {
let mut values = [50_u8; FEATURE_COUNT];
values[0] = first_feature;
LabeledSample {
sample_id: key(id),
attempt_id: key(&format!("attempt:{id}")),
speaker_id: key(speaker),
cohort_id: key("cohort"),
group_id: key(group),
clip_object: ObjectId::parse(clip).unwrap(),
primary_language: key("en"),
recording_kind: RecordingKind::Meeting,
usable_speech_ms: 1_000,
recording_quality: 90,
features: FeatureVector::new(values).unwrap(),
}
}
fn row_index(dataset: &Dataset, id: &str) -> usize {
active_rows(dataset)
.iter()
.position(|row| row.sample_id.as_ref() == id)
.unwrap()
}
fn fixture() -> (Dataset, Vec<OpenSetFold>, OpenSetFold) {
let dataset = build(
vec![
sample("d1", "d", "CLIP0007", "g-d-block", 50),
sample("c2", "c", "CLIP0005", "g-c-only", 48),
sample("b2", "b", "CLIP0004", "g-c-block", 82),
sample("a3", "a", "CLIP0007", "g-d-block", 11),
sample("c1", "c", "CLIP0004", "g-c-block", 45),
sample("a2", "a", "CLIP0002", "g-a-test", 12),
sample("b1", "b", "CLIP0003", "g-b-train", 80),
sample("a1", "a", "CLIP0001", "g-a-train", 10),
],
vec![],
)
.unwrap();
let a1 = row_index(&dataset, "a1");
let a2 = row_index(&dataset, "a2");
let a3 = row_index(&dataset, "a3");
let b1 = row_index(&dataset, "b1");
let b2 = row_index(&dataset, "b2");
let c1 = row_index(&dataset, "c1");
let c2 = row_index(&dataset, "c2");
let d1 = row_index(&dataset, "d1");
let held_out = OpenSetFold {
pseudo_unknown_speaker: key("d"),
train_indices: vec![c2, a1, b2, c1, a2, b1],
known_test_indices: vec![a3],
unknown_test_indices: vec![d1],
};
let tune_c = OpenSetFold {
pseudo_unknown_speaker: key("c"),
train_indices: vec![b1, a1],
known_test_indices: vec![b2, a2],
unknown_test_indices: vec![c2, c1],
};
let tune_b = OpenSetFold {
pseudo_unknown_speaker: key("b"),
train_indices: vec![c2, a1],
known_test_indices: vec![c1, a2],
unknown_test_indices: vec![b2, b1],
};
(dataset, vec![tune_c, tune_b], held_out)
}
fn config(absolute_threshold: f64, margin_threshold: f64) -> ModelConfig {
ModelConfig {
mask: FeatureMask::from_bits(1).unwrap(),
components: 1,
relevance: 2.0,
variance_floor: 0.01,
absolute_threshold,
margin_threshold,
}
}
fn good_config() -> ModelConfig {
config(0.0, 0.2)
}
#[test]
fn deterministic_selection_is_independent_of_input_order() {
let (dataset, mut tuning, held_out) = fixture();
let always_known = config(-1_000.0, -1_000.0);
let good = good_config();
let always_unknown = config(1_000.0, -1_000.0);
let mut candidates = vec![always_unknown.clone(), good.clone(), always_known.clone()];
let first = evaluate(EvaluationPlan {
dataset: &dataset,
tuning_folds: &tuning,
held_out_fold: &held_out,
candidates: &candidates,
})
.unwrap();
tuning.reverse();
candidates.reverse();
let second = evaluate(EvaluationPlan {
dataset: &dataset,
tuning_folds: &tuning,
held_out_fold: &held_out,
candidates: &candidates,
})
.unwrap();
assert_eq!(first, second);
assert_eq!(first.selected_config, good);
assert_eq!(
first
.candidate_evaluations
.iter()
.map(|evaluation| evaluation.config.clone())
.collect::<Vec<_>>(),
vec![always_known, good.clone(), always_unknown]
);
let selected = first
.candidate_evaluations
.iter()
.find(|evaluation| evaluation.config == good)
.unwrap();
assert_eq!(selected.tuning.known_total, 4);
assert_eq!(selected.tuning.known_correct, 4);
assert_eq!(selected.tuning.unknown_total, 4);
assert_eq!(selected.tuning.unknown_rejected, 2);
assert_eq!(first.held_out.known_total, 1);
assert_eq!(first.held_out.unknown_total, 1);
}
#[test]
fn exact_metric_ties_use_canonical_configuration_order() {
let (dataset, tuning, held_out) = fixture();
let canonical = config(-1_000.0, -1_000.0);
let later = config(-999.0, -1_000.0);
let candidates = [later, canonical.clone()];
let result = evaluate(EvaluationPlan {
dataset: &dataset,
tuning_folds: &tuning,
held_out_fold: &held_out,
candidates: &candidates,
})
.unwrap();
assert_eq!(result.selected_config, canonical);
assert_eq!(
result.candidate_evaluations[0].tuning,
result.candidate_evaluations[1].tuning
);
}
#[test]
fn unknown_rows_are_never_enrolled_and_are_counted_separately() {
let (dataset, tuning, held_out) = fixture();
let candidates = [config(1_000.0, -1_000.0)];
let result = evaluate(EvaluationPlan {
dataset: &dataset,
tuning_folds: &tuning,
held_out_fold: &held_out,
candidates: &candidates,
})
.unwrap();
let metrics = result.candidate_evaluations[0].tuning;
assert_eq!(metrics.known_total, 4);
assert_eq!(metrics.known_rejected, 4);
assert_eq!(metrics.unknown_total, 4);
assert_eq!(metrics.unknown_rejected, 4);
assert_eq!(metrics.unknown_accepted, 0);
assert_eq!(result.held_out.known_rejected, 1);
assert_eq!(result.held_out.unknown_rejected, 1);
}
fn assert_group_safe(dataset: &Dataset, fold: &OpenSetFold) {
let train_groups = fold
.train_indices
.iter()
.map(|&index| active_rows(dataset)[index].group_id.as_ref())
.collect::<BTreeSet<_>>();
let test_groups = fold
.known_test_indices
.iter()
.chain(&fold.unknown_test_indices)
.map(|&index| active_rows(dataset)[index].group_id.as_ref())
.collect::<BTreeSet<_>>();
assert!(train_groups.is_disjoint(&test_groups));
}
#[test]
fn leakage_groups_never_cross_train_and_test() {
let (dataset, tuning, held_out) = fixture();
assert_group_safe(&dataset, &held_out);
for fold in &tuning {
assert_group_safe(&dataset, fold);
}
let mut leaking = tuning
.iter()
.find(|fold| fold.pseudo_unknown_speaker.as_ref() == "c")
.unwrap()
.clone();
let c1 = row_index(&dataset, "c1");
leaking.unknown_test_indices.retain(|index| *index != c1);
leaking.train_indices.push(c1);
let candidates = [good_config()];
let folds = [leaking];
assert_eq!(
evaluate(EvaluationPlan {
dataset: &dataset,
tuning_folds: &folds,
held_out_fold: &held_out,
candidates: &candidates,
})
.unwrap_err(),
EvalError::GroupLeakage
);
}
#[test]
fn malformed_open_set_and_nested_boundaries_fail_closed() {
let (dataset, tuning, held_out) = fixture();
let tune_c = tuning
.iter()
.find(|fold| fold.pseudo_unknown_speaker.as_ref() == "c")
.unwrap();
let candidates = [good_config()];
let mut enrolled_unknown = tune_c.clone();
let c2 = row_index(&dataset, "c2");
let a1 = row_index(&dataset, "a1");
enrolled_unknown
.unknown_test_indices
.retain(|index| *index != c2);
enrolled_unknown.train_indices.push(c2);
enrolled_unknown.train_indices.retain(|index| *index != a1);
enrolled_unknown.unknown_test_indices.push(a1);
let folds = [enrolled_unknown];
assert_eq!(
evaluate(EvaluationPlan {
dataset: &dataset,
tuning_folds: &folds,
held_out_fold: &held_out,
candidates: &candidates,
})
.unwrap_err(),
EvalError::OpenSetViolation
);
let mut outside = tune_c.clone();
let a3 = row_index(&dataset, "a3");
outside.train_indices.retain(|index| *index != a1);
outside.train_indices.push(a3);
let folds = [outside];
assert_eq!(
evaluate(EvaluationPlan {
dataset: &dataset,
tuning_folds: &folds,
held_out_fold: &held_out,
candidates: &candidates,
})
.unwrap_err(),
EvalError::FoldOutsideOuterTraining
);
}
#[test]
fn malformed_plans_and_configurations_are_typed() {
let (dataset, tuning, held_out) = fixture();
let valid = good_config();
assert_eq!(
evaluate(EvaluationPlan {
dataset: &dataset,
tuning_folds: &tuning,
held_out_fold: &held_out,
candidates: &[],
})
.unwrap_err(),
EvalError::EmptyCandidates
);
assert_eq!(
evaluate(EvaluationPlan {
dataset: &dataset,
tuning_folds: &[],
held_out_fold: &held_out,
candidates: std::slice::from_ref(&valid),
})
.unwrap_err(),
EvalError::EmptyTuningFolds
);
let duplicates = [valid.clone(), valid.clone()];
assert_eq!(
evaluate(EvaluationPlan {
dataset: &dataset,
tuning_folds: &tuning,
held_out_fold: &held_out,
candidates: &duplicates,
})
.unwrap_err(),
EvalError::DuplicateCandidate
);
let duplicate_folds = [tuning[0].clone(), tuning[0].clone()];
assert_eq!(
evaluate(EvaluationPlan {
dataset: &dataset,
tuning_folds: &duplicate_folds,
held_out_fold: &held_out,
candidates: std::slice::from_ref(&valid),
})
.unwrap_err(),
EvalError::DuplicateTuningFold
);
let mut nonfinite = valid.clone();
nonfinite.absolute_threshold = f64::NAN;
assert_eq!(
evaluate(EvaluationPlan {
dataset: &dataset,
tuning_folds: &tuning,
held_out_fold: &held_out,
candidates: &[nonfinite],
})
.unwrap_err(),
EvalError::InvalidCandidate
);
let mut too_many_components = valid.clone();
too_many_components.components = 3;
assert_eq!(
evaluate(EvaluationPlan {
dataset: &dataset,
tuning_folds: &tuning,
held_out_fold: &held_out,
candidates: &[too_many_components],
})
.unwrap_err(),
EvalError::InvalidCandidate
);
let mut incomplete = held_out.clone();
incomplete.train_indices.pop();
assert_eq!(
evaluate(EvaluationPlan {
dataset: &dataset,
tuning_folds: &tuning,
held_out_fold: &incomplete,
candidates: std::slice::from_ref(&valid),
})
.unwrap_err(),
EvalError::InvalidFold
);
}
#[test]
fn dependency_contract_is_frozen_at_twenty_four_features() {
assert_eq!(FEATURE_COUNT, 24);
assert_eq!(u64::from(ALL_24), (1_u64 << FEATURE_COUNT) - 1);
assert!(FeatureVector::new([100; FEATURE_COUNT]).is_ok());
}
}