kcode_speaker_dataset/

lib.rs

use kcode_speaker_types::{Key, LabeledSample, ObjectId};
use std::collections::{BTreeMap, BTreeSet};
use std::error::Error;
use std::fmt;

/// A validated, canonically ordered pair of dataset views.
#[derive(Clone, Debug)]
pub struct Dataset {
    active: Vec<LabeledSample>,
    repeatability: Vec<LabeledSample>,
}

/// Number of deterministic known-speaker folds per pseudo-unknown speaker.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct FoldConfig {
    pub known_folds: u8,
}

/// One open-set fixture. Every index addresses [`active_rows`].
#[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>,
}

/// Distinct provider attempts for one confirmed speaker and canonical clip.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RepeatabilityGroup {
    pub speaker_id: Key,
    pub clip_object: ObjectId,
    pub sample_indices: Vec<usize>,
}

/// Validation and fold-construction failures.
#[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 {}

/// Validates and canonically orders both dataset views.
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(())
}

/// Returns the canonical active row view used by all fold indices.
pub fn active_rows(dataset: &Dataset) -> &[LabeledSample] {
    &dataset.active
}

/// Returns the canonical repeatability view used by repeatability indices.
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()
}

/// Builds deterministic leakage-safe open-set fixtures.
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)
}

/// Returns canonical clip/speaker groups having at least two distinct attempts.
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
        );
    }
}