use kcode_speaker_types::{Key, LabeledSample, ObjectId};
use std::collections::{BTreeMap, BTreeSet};
use std::error::Error;
use std::fmt;
#[derive(Clone, Debug)]
pub struct Dataset {
active: Vec<LabeledSample>,
repeatability: Vec<LabeledSample>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct FoldConfig {
pub known_folds: u8,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct OpenSetFold {
pub pseudo_unknown_speaker: Key,
pub train_indices: Vec<usize>,
pub known_test_indices: Vec<usize>,
pub unknown_test_indices: Vec<usize>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RepeatabilityGroup {
pub speaker_id: Key,
pub clip_object: ObjectId,
pub sample_indices: Vec<usize>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DatasetError {
EmptyActive,
MixedActiveCohorts,
DuplicateActiveSampleId,
DuplicateRepeatabilitySampleId,
ConflictingSharedSampleId,
DuplicateActiveObservation,
InconsistentClipGroup,
InvalidRecordingQuality,
NoUsableSpeech,
InvalidKnownFolds,
ImpossibleFold,
}
impl fmt::Display for DatasetError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let message = match self {
Self::EmptyActive => "active data is empty",
Self::MixedActiveCohorts => "active rows contain multiple cohorts",
Self::DuplicateActiveSampleId => "active rows contain a duplicate sample ID",
Self::DuplicateRepeatabilitySampleId => {
"repeatability rows contain a duplicate sample ID"
}
Self::ConflictingSharedSampleId => {
"a sample ID shared by both views has different rows"
}
Self::DuplicateActiveObservation => {
"active rows contain a duplicate clip and speaker observation"
}
Self::InconsistentClipGroup => "one active clip appears in multiple leakage groups",
Self::InvalidRecordingQuality => "recording quality exceeds 100",
Self::NoUsableSpeech => "a row has no usable speech",
Self::InvalidKnownFolds => "known_folds must be at least two",
Self::ImpossibleFold => "the requested leakage-safe folds cannot be formed",
};
f.write_str(message)
}
}
impl Error for DatasetError {}
pub fn build(
mut active_rows: Vec<LabeledSample>,
mut repeatability_rows: Vec<LabeledSample>,
) -> Result<Dataset, DatasetError> {
if active_rows.is_empty() {
return Err(DatasetError::EmptyActive);
}
active_rows.sort_by(|a, b| a.sample_id.as_ref().cmp(b.sample_id.as_ref()));
repeatability_rows.sort_by(|a, b| a.sample_id.as_ref().cmp(b.sample_id.as_ref()));
validate_rows(&active_rows)?;
validate_rows(&repeatability_rows)?;
if active_rows
.windows(2)
.any(|rows| rows[0].sample_id == rows[1].sample_id)
{
return Err(DatasetError::DuplicateActiveSampleId);
}
if repeatability_rows
.windows(2)
.any(|rows| rows[0].sample_id == rows[1].sample_id)
{
return Err(DatasetError::DuplicateRepeatabilitySampleId);
}
let cohort = active_rows[0].cohort_id.as_ref();
if active_rows
.iter()
.any(|row| row.cohort_id.as_ref() != cohort)
{
return Err(DatasetError::MixedActiveCohorts);
}
let mut observations = BTreeSet::new();
let mut clip_groups: BTreeMap<&str, &str> = BTreeMap::new();
for row in &active_rows {
if !observations.insert((row.clip_object.as_ref(), row.speaker_id.as_ref())) {
return Err(DatasetError::DuplicateActiveObservation);
}
match clip_groups.get(row.clip_object.as_ref()) {
Some(group) if *group != row.group_id.as_ref() => {
return Err(DatasetError::InconsistentClipGroup);
}
Some(_) => {}
None => {
clip_groups.insert(row.clip_object.as_ref(), row.group_id.as_ref());
}
}
}
let active_by_id: BTreeMap<&str, &LabeledSample> = active_rows
.iter()
.map(|row| (row.sample_id.as_ref(), row))
.collect();
for row in &repeatability_rows {
if let Some(active) = active_by_id.get(row.sample_id.as_ref())
&& *active != row
{
return Err(DatasetError::ConflictingSharedSampleId);
}
}
Ok(Dataset {
active: active_rows,
repeatability: repeatability_rows,
})
}
fn validate_rows(rows: &[LabeledSample]) -> Result<(), DatasetError> {
for row in rows {
if row.recording_quality > 100 {
return Err(DatasetError::InvalidRecordingQuality);
}
if row.usable_speech_ms == 0 {
return Err(DatasetError::NoUsableSpeech);
}
}
Ok(())
}
pub fn active_rows(dataset: &Dataset) -> &[LabeledSample] {
&dataset.active
}
pub fn repeatability_rows(dataset: &Dataset) -> &[LabeledSample] {
&dataset.repeatability
}
#[derive(Debug)]
struct Group {
id: String,
indices: Vec<usize>,
speaker_counts: BTreeMap<String, usize>,
}
fn active_groups(dataset: &Dataset) -> Vec<Group> {
let mut indices_by_group: BTreeMap<String, Vec<usize>> = BTreeMap::new();
for (index, row) in dataset.active.iter().enumerate() {
indices_by_group
.entry(row.group_id.as_ref().to_owned())
.or_default()
.push(index);
}
indices_by_group
.into_iter()
.map(|(id, indices)| {
let mut speaker_counts = BTreeMap::new();
for &index in &indices {
*speaker_counts
.entry(dataset.active[index].speaker_id.as_ref().to_owned())
.or_default() += 1;
}
Group {
id,
indices,
speaker_counts,
}
})
.collect()
}
pub fn open_set_folds(
dataset: &Dataset,
config: FoldConfig,
) -> Result<Vec<OpenSetFold>, DatasetError> {
let fold_count = usize::from(config.known_folds);
if fold_count < 2 {
return Err(DatasetError::InvalidKnownFolds);
}
let groups = active_groups(dataset);
let speakers: BTreeMap<String, Key> = dataset
.active
.iter()
.map(|row| (row.speaker_id.as_ref().to_owned(), row.speaker_id.clone()))
.collect();
let mut output = Vec::with_capacity(speakers.len() * fold_count);
for (pseudo_name, pseudo_key) in speakers {
let (blocked, mut eligible): (Vec<&Group>, Vec<&Group>) = groups
.iter()
.partition(|group| group.speaker_counts.contains_key(&pseudo_name));
if blocked.is_empty() || eligible.len() < fold_count {
return Err(DatasetError::ImpossibleFold);
}
eligible.sort_by(|a, b| {
b.indices
.len()
.cmp(&a.indices.len())
.then_with(|| b.speaker_counts.len().cmp(&a.speaker_counts.len()))
.then_with(|| a.id.cmp(&b.id))
});
let mut test_counts = vec![BTreeMap::<String, usize>::new(); fold_count];
let mut row_counts = vec![0_usize; fold_count];
let mut group_counts = vec![0_usize; fold_count];
let mut assignment = BTreeMap::<&str, usize>::new();
for group in &eligible {
let chosen = (0..fold_count)
.min_by_key(|&fold| {
let speaker_cost: usize = group
.speaker_counts
.iter()
.map(|(speaker, addition)| {
let current =
test_counts[fold].get(speaker).copied().unwrap_or_default();
2 * current * addition + addition * addition
})
.sum();
(speaker_cost, row_counts[fold], group_counts[fold], fold)
})
.expect("fold count is nonzero");
assignment.insert(group.id.as_str(), chosen);
for (speaker, count) in &group.speaker_counts {
*test_counts[chosen].entry(speaker.clone()).or_default() += count;
}
row_counts[chosen] += group.indices.len();
group_counts[chosen] += 1;
}
if group_counts.contains(&0) {
return Err(DatasetError::ImpossibleFold);
}
for fold in 0..fold_count {
let mut train_indices: Vec<usize> = Vec::new();
let mut known_test_indices: Vec<usize> = Vec::new();
let mut unknown_test_indices: Vec<usize> = Vec::new();
for group in &blocked {
for &index in &group.indices {
if dataset.active[index].speaker_id.as_ref() == pseudo_name {
unknown_test_indices.push(index);
} else {
known_test_indices.push(index);
}
}
}
for group in &eligible {
if assignment[group.id.as_str()] == fold {
known_test_indices.extend(&group.indices);
} else {
train_indices.extend(&group.indices);
}
}
train_indices.sort_unstable();
known_test_indices.sort_unstable();
unknown_test_indices.sort_unstable();
if train_indices.is_empty()
|| known_test_indices.is_empty()
|| unknown_test_indices.is_empty()
{
return Err(DatasetError::ImpossibleFold);
}
let trained_speakers: BTreeSet<&str> = train_indices
.iter()
.map(|index: &usize| dataset.active[*index].speaker_id.as_ref())
.collect();
if known_test_indices.iter().any(|index: &usize| {
!trained_speakers.contains(dataset.active[*index].speaker_id.as_ref())
}) {
return Err(DatasetError::ImpossibleFold);
}
output.push(OpenSetFold {
pseudo_unknown_speaker: pseudo_key.clone(),
train_indices,
known_test_indices,
unknown_test_indices,
});
}
}
Ok(output)
}
pub fn repeatability_groups(dataset: &Dataset) -> Vec<RepeatabilityGroup> {
let mut grouped: BTreeMap<(String, String), Vec<usize>> = BTreeMap::new();
for (index, row) in dataset.repeatability.iter().enumerate() {
grouped
.entry((
row.speaker_id.as_ref().to_owned(),
row.clip_object.as_ref().to_owned(),
))
.or_default()
.push(index);
}
grouped
.into_values()
.filter(|indices| {
indices
.iter()
.map(|&index| dataset.repeatability[index].attempt_id.as_ref())
.collect::<BTreeSet<_>>()
.len()
>= 2
})
.map(|sample_indices| {
let row = &dataset.repeatability[sample_indices[0]];
RepeatabilityGroup {
speaker_id: row.speaker_id.clone(),
clip_object: row.clip_object.clone(),
sample_indices,
}
})
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
use kcode_speaker_types::{FEATURE_COUNT, FeatureVector, RecordingKind};
fn key(value: &str) -> Key {
Key::parse(value).unwrap()
}
fn object(value: &str) -> ObjectId {
ObjectId::parse(value).unwrap()
}
fn sample_with_attempt(
id: &str,
attempt: &str,
speaker: &str,
clip: &str,
group: &str,
) -> LabeledSample {
LabeledSample {
sample_id: key(id),
attempt_id: key(attempt),
speaker_id: key(speaker),
cohort_id: key("cohort"),
group_id: key(group),
clip_object: object(clip),
recording_kind: RecordingKind::Meeting,
recording_quality: 100,
usable_speech_ms: 900,
primary_language: key("en"),
features: FeatureVector::new([50; FEATURE_COUNT]).unwrap(),
}
}
fn sample(id: &str, speaker: &str, clip: &str, group: &str) -> LabeledSample {
sample_with_attempt(id, id, speaker, clip, group)
}
fn fixture_rows() -> Vec<LabeledSample> {
vec![
sample("a1", "a", "CLIP0001", "a-source"),
sample("a2", "a", "CLIP0002", "a-source"),
sample("a3", "a", "CLIP0003", "a-three"),
sample("a4", "a", "CLIP0004", "meeting"),
sample("b1", "b", "CLIP0011", "b-one"),
sample("b2", "b", "CLIP0012", "b-two"),
sample("b3", "b", "CLIP0013", "b-three"),
sample("b4", "b", "CLIP0004", "meeting"),
sample("c1", "c", "CLIP0021", "c-source"),
sample("c2", "c", "CLIP0022", "c-source"),
sample("c3", "c", "CLIP0023", "c-three"),
]
}
fn repeat_rows() -> Vec<LabeledSample> {
vec![
sample_with_attempt("r3", "try-3", "a", "CLIP0031", "repeat"),
sample_with_attempt("r1", "try-1", "a", "CLIP0030", "repeat"),
sample_with_attempt("r2", "try-2", "a", "CLIP0030", "repeat"),
sample_with_attempt("r4", "try-4", "b", "CLIP0032", "repeat"),
sample_with_attempt("r5", "same-try", "c", "CLIP0033", "repeat"),
sample_with_attempt("r6", "same-try", "c", "CLIP0033", "repeat"),
]
}
fn side(fold: &OpenSetFold, index: usize) -> u8 {
if fold.train_indices.contains(&index) {
1
} else if fold.known_test_indices.contains(&index)
|| fold.unknown_test_indices.contains(&index)
{
2
} else {
0
}
}
#[test]
fn groups_never_cross_and_co_speakers_are_test_only() {
let dataset = build(fixture_rows(), vec![]).unwrap();
let folds = open_set_folds(&dataset, FoldConfig { known_folds: 2 }).unwrap();
for fold in &folds {
let mut group_sides: BTreeMap<&str, u8> = BTreeMap::new();
for (index, row) in active_rows(&dataset).iter().enumerate() {
let current = side(fold, index);
assert_ne!(current, 0);
let previous = group_sides.entry(row.group_id.as_ref()).or_insert(current);
assert_eq!(*previous, current);
}
}
let meeting_b = active_rows(&dataset)
.iter()
.position(|row| row.sample_id.as_ref() == "b4")
.unwrap();
for fold in folds
.iter()
.filter(|fold| fold.pseudo_unknown_speaker.as_ref() == "a")
{
assert!(fold.known_test_indices.contains(&meeting_b));
assert!(!fold.train_indices.contains(&meeting_b));
}
}
#[test]
fn sibling_duplicate_excerpt_and_overlap_group_stays_whole() {
let dataset = build(fixture_rows(), vec![]).unwrap();
let folds = open_set_folds(&dataset, FoldConfig { known_folds: 2 }).unwrap();
let related: Vec<_> = active_rows(&dataset)
.iter()
.enumerate()
.filter(|(_, row)| row.group_id.as_ref() == "a-source")
.map(|(index, _)| index)
.collect();
assert_eq!(related.len(), 2);
for fold in folds {
assert_eq!(side(&fold, related[0]), side(&fold, related[1]));
}
}
#[test]
fn permutation_balancing_and_index_views_are_stable() {
let active = fixture_rows();
let repeatability = repeat_rows();
let mut reversed_active = active.clone();
let mut reversed_repeatability = repeatability.clone();
reversed_active.reverse();
reversed_repeatability.reverse();
let first = build(active, repeatability).unwrap();
let second = build(reversed_active, reversed_repeatability).unwrap();
assert_eq!(
active_rows(&first)
.iter()
.map(|row| row.sample_id.as_ref())
.collect::<Vec<_>>(),
active_rows(&second)
.iter()
.map(|row| row.sample_id.as_ref())
.collect::<Vec<_>>()
);
let first_folds = open_set_folds(&first, FoldConfig { known_folds: 2 }).unwrap();
assert_eq!(
first_folds,
open_set_folds(&second, FoldConfig { known_folds: 2 }).unwrap()
);
assert_eq!(repeatability_groups(&first), repeatability_groups(&second));
assert_eq!(first_folds.len(), 6);
let b_folds: Vec<_> = first_folds
.iter()
.filter(|fold| fold.pseudo_unknown_speaker.as_ref() == "b")
.collect();
let known_ids = |fold: &OpenSetFold| {
fold.known_test_indices
.iter()
.map(|&index| active_rows(&first)[index].sample_id.as_ref())
.collect::<BTreeSet<_>>()
};
assert_eq!(
known_ids(b_folds[0]),
BTreeSet::from(["a1", "a2", "a4", "c3"])
);
assert_eq!(
known_ids(b_folds[1]),
BTreeSet::from(["a3", "a4", "c1", "c2"])
);
for fold in first_folds {
for index in fold
.train_indices
.iter()
.chain(&fold.known_test_indices)
.chain(&fold.unknown_test_indices)
{
assert!(active_rows(&first).get(*index).is_some());
}
}
for group in repeatability_groups(&first) {
for index in group.sample_indices {
assert!(repeatability_rows(&first).get(index).is_some());
}
}
}
#[test]
fn exact_repeatability_grouping_requires_distinct_attempts() {
let dataset = build(fixture_rows(), repeat_rows()).unwrap();
let groups = repeatability_groups(&dataset);
assert_eq!(groups.len(), 1);
assert_eq!(groups[0].speaker_id.as_ref(), "a");
assert_eq!(groups[0].clip_object.as_ref(), "CLIP0030");
let ids: Vec<_> = groups[0]
.sample_indices
.iter()
.map(|&index| repeatability_rows(&dataset)[index].sample_id.as_ref())
.collect();
assert_eq!(ids, ["r1", "r2"]);
}
#[test]
fn duplicate_and_cross_view_rules_are_exact() {
let row = sample("same", "a", "CLIP0040", "g");
assert_eq!(
build(vec![row.clone(), row.clone()], vec![]).unwrap_err(),
DatasetError::DuplicateActiveSampleId
);
assert_eq!(
build(fixture_rows(), vec![row.clone(), row.clone()]).unwrap_err(),
DatasetError::DuplicateRepeatabilitySampleId
);
let dataset = build(vec![row.clone()], vec![row.clone()]).unwrap();
assert_eq!(active_rows(&dataset)[0], repeatability_rows(&dataset)[0]);
let mut changed = row.clone();
changed.usable_speech_ms = 1;
assert_eq!(
build(vec![row], vec![changed]).unwrap_err(),
DatasetError::ConflictingSharedSampleId
);
}
#[test]
fn sparse_and_impossible_configs_fail_closed() {
let dataset = build(fixture_rows(), vec![]).unwrap();
assert_eq!(
open_set_folds(&dataset, FoldConfig { known_folds: 1 }).unwrap_err(),
DatasetError::InvalidKnownFolds
);
assert_eq!(
open_set_folds(&dataset, FoldConfig { known_folds: 20 }).unwrap_err(),
DatasetError::ImpossibleFold
);
let rows = vec![
sample("a1", "a", "CLIP0050", "ab"),
sample("b1", "b", "CLIP0050", "ab"),
sample("b2", "b", "CLIP0051", "b-only"),
sample("c1", "c", "CLIP0052", "c-one"),
sample("c2", "c", "CLIP0053", "c-two"),
];
let sparse = build(rows, vec![]).unwrap();
assert_eq!(
open_set_folds(&sparse, FoldConfig { known_folds: 2 }).unwrap_err(),
DatasetError::ImpossibleFold
);
}
#[test]
fn active_observation_and_supplied_clip_group_are_validated() {
let first = sample("x1", "a", "CLIP0060", "one");
let duplicate = sample("x2", "a", "CLIP0060", "one");
assert_eq!(
build(vec![first.clone(), duplicate], vec![]).unwrap_err(),
DatasetError::DuplicateActiveObservation
);
let other_speaker = sample("x3", "b", "CLIP0060", "two");
assert_eq!(
build(vec![first, other_speaker], vec![]).unwrap_err(),
DatasetError::InconsistentClipGroup
);
}
#[test]
fn cohort_quality_and_usable_speech_are_validated() {
let mut mixed = sample("x2", "b", "CLIP0062", "two");
mixed.cohort_id = key("other");
assert_eq!(
build(vec![sample("x1", "a", "CLIP0061", "one"), mixed], vec![]).unwrap_err(),
DatasetError::MixedActiveCohorts
);
let mut poor = sample("x3", "a", "CLIP0063", "three");
poor.recording_quality = 101;
assert_eq!(
build(vec![poor], vec![]).unwrap_err(),
DatasetError::InvalidRecordingQuality
);
let mut silent = sample("x4", "a", "CLIP0064", "four");
silent.usable_speech_ms = 0;
assert_eq!(
build(vec![silent], vec![]).unwrap_err(),
DatasetError::NoUsableSpeech
);
}
}