kcode_speaker_model/

lib.rs

#![forbid(unsafe_code)]

use kcode_diag_gmm::{
    DiagGmm, FitConfig as GmmFitConfig, fit as fit_gmm, log_likelihood, means, responsibilities,
    with_means,
};
use kcode_speaker_types::{
    FEATURE_COUNT, FeatureMask, FeatureVector, Key, LabeledSample, RecordingKind,
};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::fmt;

const SNAPSHOT_VERSION: u8 = 2;
const WIRE_FEATURE_COUNT: u8 = FEATURE_COUNT as u8;
const MAX_ITERATIONS: u16 = 200;
const RELATIVE_TOLERANCE: f64 = 1e-8;

pub const ALL_24: FeatureMask = match FeatureMask::from_bits((1_u64 << FEATURE_COUNT) - 1) {
    Ok(mask) => mask,
    Err(_) => panic!("FEATURE_COUNT must define a valid nonempty feature mask"),
};

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ModelConfig {
    pub mask: FeatureMask,
    pub components: u8,
    pub relevance: f64,
    pub variance_floor: f64,
    pub absolute_threshold: f64,
    pub margin_threshold: f64,
}

pub struct FitInput<'a> {
    pub cohort_id: &'a Key,
    pub samples: &'a [LabeledSample],
    pub config: ModelConfig,
}

/// The complete model input needed for one labelled feature row.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct TrainingRow {
    pub sample_id: Key,
    pub speaker_id: Key,
    pub cohort_id: Key,
    pub features: FeatureVector,
}

/// A minimal fitting request for callers that do not own media provenance.
pub struct FitRowsInput<'a> {
    pub cohort_id: &'a Key,
    pub rows: &'a [TrainingRow],
    pub config: ModelConfig,
}

#[derive(Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ModelSnapshot {
    version: u8,
    feature_count: u8,
    cohort_id: Key,
    config: ModelConfig,
    selected_indices: Vec<u8>,
    normalizer_mean: Vec<f64>,
    normalizer_std: Vec<f64>,
    ubm: DiagGmm,
    speakers: Vec<SpeakerModel>,
    training_sample_count: usize,
    manifest_sha256: [u8; 32],
}

#[derive(Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct SpeakerModel {
    speaker_id: Key,
    sample_count: usize,
    occupancy: Vec<f64>,
    first_moments: Vec<Vec<f64>>,
    adapted_means: Vec<Vec<f64>>,
}

#[derive(Clone, Debug, PartialEq)]
pub struct CandidateScore {
    pub speaker_id: Key,
    pub llr: f64,
}

#[derive(Clone, Debug, PartialEq)]
pub enum Decision {
    Known { speaker_id: Key },
    Unknown,
}

#[derive(Clone, Debug, PartialEq)]
pub struct Identification {
    pub decision: Decision,
    pub best: CandidateScore,
    pub runner_up: Option<CandidateScore>,
    pub absolute_pass: bool,
    pub margin_pass: bool,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ModelError {
    InvalidConfig,
    InvalidSamples,
    CohortMismatch,
    DuplicateSampleId,
    ZeroVariance,
    Nonconverged,
    Numerical,
    Gmm,
    Serialization,
    MalformedArtifact,
}

impl fmt::Display for ModelError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        let text = match self {
            Self::InvalidConfig => "invalid model configuration",
            Self::InvalidSamples => "invalid training samples",
            Self::CohortMismatch => "cohort does not match",
            Self::DuplicateSampleId => "duplicate training sample ID",
            Self::ZeroVariance => "selected training feature has zero variance",
            Self::Nonconverged => "diagonal GMM did not converge",
            Self::Numerical => "nonfinite model computation",
            Self::Gmm => "diagonal GMM operation failed",
            Self::Serialization => "snapshot serialization failed",
            Self::MalformedArtifact => "malformed or incompatible snapshot artifact",
        };
        formatter.write_str(text)
    }
}

impl std::error::Error for ModelError {}

pub fn fit(input: FitInput<'_>) -> Result<ModelSnapshot, ModelError> {
    if input.samples.is_empty() {
        return Err(ModelError::InvalidSamples);
    }
    validate_config(&input.config, input.samples.len())?;

    let mut samples = input.samples.iter().collect::<Vec<_>>();
    samples.sort_by(|left, right| left.sample_id.cmp(&right.sample_id));

    if samples
        .iter()
        .any(|sample| sample.cohort_id != *input.cohort_id)
    {
        return Err(ModelError::CohortMismatch);
    }
    if samples
        .windows(2)
        .any(|pair| pair[0].sample_id == pair[1].sample_id)
    {
        return Err(ModelError::DuplicateSampleId);
    }

    let digest = manifest(&samples);
    let rows = samples
        .into_iter()
        .map(|sample| ModelRow {
            speaker_id: &sample.speaker_id,
            features: &sample.features,
        })
        .collect();
    fit_rows_inner(input.cohort_id, rows, input.config, digest)
}

/// Fits the same deterministic model from only IDs and feature vectors.
pub fn fit_rows(input: FitRowsInput<'_>) -> Result<ModelSnapshot, ModelError> {
    if input.rows.is_empty() {
        return Err(ModelError::InvalidSamples);
    }
    validate_config(&input.config, input.rows.len())?;

    let mut rows = input.rows.iter().collect::<Vec<_>>();
    rows.sort_by(|left, right| left.sample_id.cmp(&right.sample_id));
    if rows.iter().any(|row| row.cohort_id != *input.cohort_id) {
        return Err(ModelError::CohortMismatch);
    }
    if rows
        .windows(2)
        .any(|pair| pair[0].sample_id == pair[1].sample_id)
    {
        return Err(ModelError::DuplicateSampleId);
    }

    let digest = row_manifest(&rows);
    let rows = rows
        .into_iter()
        .map(|row| ModelRow {
            speaker_id: &row.speaker_id,
            features: &row.features,
        })
        .collect();
    fit_rows_inner(input.cohort_id, rows, input.config, digest)
}

struct ModelRow<'a> {
    speaker_id: &'a Key,
    features: &'a FeatureVector,
}

fn fit_rows_inner(
    cohort_id: &Key,
    samples: Vec<ModelRow<'_>>,
    config: ModelConfig,
    manifest_sha256: [u8; 32],
) -> Result<ModelSnapshot, ModelError> {
    let selected = mask_indices(config.mask);
    let raw_rows = samples
        .iter()
        .map(|sample| {
            selected
                .iter()
                .map(|index| f64::from(sample.features.as_ref()[usize::from(*index)]))
                .collect::<Vec<_>>()
        })
        .collect::<Vec<_>>();
    let (normalizer_mean, normalizer_std) = fit_normalizer(&raw_rows)?;
    let rows = raw_rows
        .iter()
        .map(|row| normalize(row, &normalizer_mean, &normalizer_std))
        .collect::<Result<Vec<_>, _>>()?;

    let ubm = fit_ubm_with_limits(
        &rows,
        config.components,
        config.variance_floor,
        MAX_ITERATIONS,
        RELATIVE_TOLERANCE,
    )?;
    let ubm_means = means(&ubm);

    let mut speaker_ids = samples
        .iter()
        .map(|sample| (*sample.speaker_id).clone())
        .collect::<Vec<_>>();
    speaker_ids.sort();
    speaker_ids.dedup();
    if speaker_ids.is_empty() {
        return Err(ModelError::InvalidSamples);
    }

    let components = usize::from(config.components);
    let dimension = selected.len();
    let mut speakers = speaker_ids
        .into_iter()
        .map(|speaker_id| SpeakerModel {
            speaker_id,
            sample_count: 0,
            occupancy: vec![0.0; components],
            first_moments: vec![vec![0.0; dimension]; components],
            adapted_means: Vec::new(),
        })
        .collect::<Vec<_>>();

    for (sample, row) in samples.iter().zip(&rows) {
        let speaker_index = speakers
            .binary_search_by(|speaker| speaker.speaker_id.cmp(sample.speaker_id))
            .map_err(|_| ModelError::InvalidSamples)?;
        let gamma = responsibilities(&ubm, row).map_err(|_| ModelError::Gmm)?;
        if gamma.len() != components {
            return Err(ModelError::Numerical);
        }

        let speaker = &mut speakers[speaker_index];
        speaker.sample_count += 1;
        for (component, responsibility) in gamma.iter().copied().enumerate() {
            if !responsibility.is_finite() || responsibility < 0.0 {
                return Err(ModelError::Numerical);
            }

            speaker.occupancy[component] += responsibility;
            if !speaker.occupancy[component].is_finite() {
                return Err(ModelError::Numerical);
            }
            for (sum, value) in speaker.first_moments[component].iter_mut().zip(row) {
                *sum += responsibility * value;
                if !sum.is_finite() {
                    return Err(ModelError::Numerical);
                }
            }
        }
    }

    for speaker in &mut speakers {
        speaker.adapted_means = map_means(
            ubm_means,
            &speaker.occupancy,
            &speaker.first_moments,
            config.relevance,
        )?;
    }

    let model = ModelSnapshot {
        version: SNAPSHOT_VERSION,
        feature_count: WIRE_FEATURE_COUNT,
        cohort_id: cohort_id.clone(),
        config,
        selected_indices: selected,
        normalizer_mean,
        normalizer_std,
        ubm,
        speakers,
        training_sample_count: samples.len(),
        manifest_sha256,
    };
    validate_snapshot(&model)?;
    Ok(model)
}

pub fn identify(
    model: &ModelSnapshot,
    cohort_id: &Key,
    features: &FeatureVector,
) -> Result<Identification, ModelError> {
    validate_snapshot(model)?;
    if cohort_id != &model.cohort_id {
        return Err(ModelError::CohortMismatch);
    }

    let raw = model
        .selected_indices
        .iter()
        .map(|index| f64::from(features.as_ref()[usize::from(*index)]))
        .collect::<Vec<_>>();
    let row = normalize(&raw, &model.normalizer_mean, &model.normalizer_std)?;
    let ubm_score = log_likelihood(&model.ubm, &row).map_err(|_| ModelError::Gmm)?;
    if !ubm_score.is_finite() {
        return Err(ModelError::Numerical);
    }

    let mut scores = Vec::with_capacity(model.speakers.len());
    for speaker in &model.speakers {
        let adapted =
            with_means(&model.ubm, speaker.adapted_means.clone()).map_err(|_| ModelError::Gmm)?;
        let score = log_likelihood(&adapted, &row).map_err(|_| ModelError::Gmm)? - ubm_score;
        if !score.is_finite() {
            return Err(ModelError::Numerical);
        }
        scores.push(CandidateScore {
            speaker_id: speaker.speaker_id.clone(),
            llr: score,
        });
    }
    scores.sort_by(|left, right| {
        right
            .llr
            .total_cmp(&left.llr)
            .then_with(|| left.speaker_id.cmp(&right.speaker_id))
    });

    let mut ranked = scores.into_iter();
    let best = ranked.next().ok_or(ModelError::MalformedArtifact)?;
    let runner_up = ranked.next();
    let absolute_pass = best.llr >= model.config.absolute_threshold;
    let margin_pass = runner_up
        .as_ref()
        .is_none_or(|runner| best.llr - runner.llr >= model.config.margin_threshold);
    let decision = if absolute_pass && margin_pass {
        Decision::Known {
            speaker_id: best.speaker_id.clone(),
        }
    } else {
        Decision::Unknown
    };

    Ok(Identification {
        decision,
        best,
        runner_up,
        absolute_pass,
        margin_pass,
    })
}

pub fn encode(model: &ModelSnapshot) -> Result<Vec<u8>, ModelError> {
    validate_snapshot(model)?;
    serde_json::to_vec(model).map_err(|_| ModelError::Serialization)
}

pub fn decode(bytes: &[u8]) -> Result<ModelSnapshot, ModelError> {
    let model = serde_json::from_slice::<ModelSnapshot>(bytes)
        .map_err(|_| ModelError::MalformedArtifact)?;
    validate_snapshot(&model).map_err(|_| ModelError::MalformedArtifact)?;
    Ok(model)
}

fn validate_config(config: &ModelConfig, sample_count: usize) -> Result<(), ModelError> {
    if config.components == 0
        || usize::from(config.components) > sample_count
        || !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(ModelError::InvalidConfig);
    }
    Ok(())
}

fn mask_indices(mask: FeatureMask) -> Vec<u8> {
    let bits = u64::from(mask);
    (0..FEATURE_COUNT)
        .filter(|index| bits & (1_u64 << index) != 0)
        .map(|index| index as u8)
        .collect()
}

fn fit_normalizer(rows: &[Vec<f64>]) -> Result<(Vec<f64>, Vec<f64>), ModelError> {
    let dimension = rows.first().map_or(0, Vec::len);
    if rows.is_empty() || dimension == 0 || rows.iter().any(|row| row.len() != dimension) {
        return Err(ModelError::InvalidSamples);
    }

    let count = rows.len() as f64;
    let mut mean = vec![0.0; dimension];
    for row in rows {
        for (sum, value) in mean.iter_mut().zip(row) {
            *sum += value;
        }
    }
    for value in &mut mean {
        *value /= count;
    }
    if mean.iter().any(|value| !value.is_finite()) {
        return Err(ModelError::Numerical);
    }

    let mut standard_deviation = vec![0.0; dimension];
    for row in rows {
        for (sum, (value, center)) in standard_deviation.iter_mut().zip(row.iter().zip(&mean)) {
            let delta = value - center;
            *sum += delta * delta;
        }
    }
    for value in &mut standard_deviation {
        *value = (*value / count).sqrt();
        if !value.is_finite() || *value <= 0.0 {
            return Err(ModelError::ZeroVariance);
        }
    }

    Ok((mean, standard_deviation))
}

fn normalize(
    row: &[f64],
    mean: &[f64],
    standard_deviation: &[f64],
) -> Result<Vec<f64>, ModelError> {
    if row.len() != mean.len() || mean.len() != standard_deviation.len() {
        return Err(ModelError::Numerical);
    }

    row.iter()
        .zip(mean.iter().zip(standard_deviation))
        .map(|(value, (center, scale))| {
            let normalized = (value - center) / scale;
            normalized
                .is_finite()
                .then_some(normalized)
                .ok_or(ModelError::Numerical)
        })
        .collect()
}

fn fit_ubm_with_limits(
    rows: &[Vec<f64>],
    components: u8,
    variance_floor: f64,
    max_iterations: u16,
    relative_tolerance: f64,
) -> Result<DiagGmm, ModelError> {
    let result = fit_gmm(
        rows,
        GmmFitConfig {
            components,
            max_iterations,
            relative_tolerance,
            variance_floor,
        },
    )
    .map_err(|_| ModelError::Gmm)?;
    if !result.converged {
        return Err(ModelError::Nonconverged);
    }
    Ok(result.model)
}

fn map_means(
    ubm_means: &[Vec<f64>],
    occupancy: &[f64],
    first_moments: &[Vec<f64>],
    relevance: f64,
) -> Result<Vec<Vec<f64>>, ModelError> {
    if occupancy.len() != ubm_means.len()
        || first_moments.len() != ubm_means.len()
        || !relevance.is_finite()
        || relevance <= 0.0
    {
        return Err(ModelError::Numerical);
    }

    let mut adapted = Vec::with_capacity(ubm_means.len());
    for ((background, mass), first) in ubm_means
        .iter()
        .zip(occupancy.iter().copied())
        .zip(first_moments)
    {
        if !mass.is_finite()
            || mass < 0.0
            || first.len() != background.len()
            || first.iter().any(|value| !value.is_finite())
            || mass == 0.0 && first.iter().any(|value| *value != 0.0)
        {
            return Err(ModelError::Numerical);
        }

        let mut component = background.clone();
        if mass > 0.0 {
            let alpha = mass / (mass + relevance);
            if !alpha.is_finite() {
                return Err(ModelError::Numerical);
            }
            for ((value, first_value), background_value) in
                component.iter_mut().zip(first).zip(background)
            {
                let empirical = first_value / mass;
                *value = alpha * empirical + (1.0 - alpha) * background_value;
                if !value.is_finite() {
                    return Err(ModelError::Numerical);
                }
            }
        }
        adapted.push(component);
    }
    Ok(adapted)
}

fn manifest(samples: &[&LabeledSample]) -> [u8; 32] {
    let mut ordered = samples.to_vec();
    ordered.sort_by(|left, right| left.sample_id.cmp(&right.sample_id));

    let mut digest = Sha256::new();
    manifest_field(&mut digest, 0, b"kcode-speaker-model/training-manifest/2");
    digest.update((ordered.len() as u64).to_be_bytes());

    for sample in ordered {
        digest.update([0xfe]);
        manifest_field(&mut digest, 1, sample.sample_id.as_ref().as_bytes());
        manifest_field(&mut digest, 2, sample.attempt_id.as_ref().as_bytes());
        manifest_field(&mut digest, 3, sample.speaker_id.as_ref().as_bytes());
        manifest_field(&mut digest, 4, sample.cohort_id.as_ref().as_bytes());
        manifest_field(&mut digest, 5, sample.group_id.as_ref().as_bytes());
        manifest_field(&mut digest, 6, sample.clip_object.as_ref().as_bytes());
        manifest_field(&mut digest, 7, sample.primary_language.as_ref().as_bytes());
        manifest_field(
            &mut digest,
            8,
            &[recording_kind_byte(&sample.recording_kind)],
        );
        manifest_field(&mut digest, 9, &sample.usable_speech_ms.to_be_bytes());
        manifest_field(&mut digest, 10, &[sample.recording_quality]);
        manifest_field(&mut digest, 11, sample.features.as_ref());
        digest.update([0xff]);
    }

    digest.finalize().into()
}

fn row_manifest(rows: &[&TrainingRow]) -> [u8; 32] {
    let mut digest = Sha256::new();
    manifest_field(
        &mut digest,
        0,
        b"kcode-speaker-model/training-row-manifest/1",
    );
    digest.update((rows.len() as u64).to_be_bytes());
    for row in rows {
        digest.update([0xfe]);
        manifest_field(&mut digest, 1, row.sample_id.as_ref().as_bytes());
        manifest_field(&mut digest, 2, row.speaker_id.as_ref().as_bytes());
        manifest_field(&mut digest, 3, row.cohort_id.as_ref().as_bytes());
        manifest_field(&mut digest, 4, row.features.as_ref());
        digest.update([0xff]);
    }
    digest.finalize().into()
}

fn manifest_field(digest: &mut Sha256, tag: u8, bytes: &[u8]) {
    digest.update([tag]);
    digest.update((bytes.len() as u64).to_be_bytes());
    digest.update(bytes);
}

fn recording_kind_byte(kind: &RecordingKind) -> u8 {
    match kind {
        RecordingKind::VoiceNote => 0,
        RecordingKind::Meeting => 1,
        RecordingKind::Call => 2,
        RecordingKind::Other => 3,
    }
}

fn validate_snapshot(model: &ModelSnapshot) -> Result<(), ModelError> {
    if model.version != SNAPSHOT_VERSION
        || model.feature_count != WIRE_FEATURE_COUNT
        || model.training_sample_count == 0
    {
        return Err(ModelError::MalformedArtifact);
    }
    validate_config(&model.config, model.training_sample_count)
        .map_err(|_| ModelError::MalformedArtifact)?;

    let expected_indices = mask_indices(model.config.mask);
    if model.selected_indices != expected_indices
        || model.normalizer_mean.len() != expected_indices.len()
        || model.normalizer_std.len() != expected_indices.len()
        || model.normalizer_mean.iter().any(|value| !value.is_finite())
        || model
            .normalizer_std
            .iter()
            .any(|value| !value.is_finite() || *value <= 0.0)
    {
        return Err(ModelError::MalformedArtifact);
    }

    let dimension = expected_indices.len();
    let components = usize::from(model.config.components);
    let ubm_means = means(&model.ubm);
    let zero_row = vec![0.0; dimension];
    if ubm_means.len() != components
        || ubm_means
            .iter()
            .any(|row| row.len() != dimension || row.iter().any(|value| !value.is_finite()))
        || log_likelihood(&model.ubm, &zero_row).is_err()
        || model.speakers.is_empty()
        || model.speakers.len() > model.training_sample_count
    {
        return Err(ModelError::MalformedArtifact);
    }

    let mut counted_samples = 0_usize;
    for (index, speaker) in model.speakers.iter().enumerate() {
        let out_of_order = index > 0
            && model.speakers[index - 1].speaker_id.as_ref() >= speaker.speaker_id.as_ref();
        if speaker.sample_count == 0
            || out_of_order
            || speaker.occupancy.len() != components
            || speaker.first_moments.len() != components
            || speaker.adapted_means.len() != components
        {
            return Err(ModelError::MalformedArtifact);
        }

        counted_samples = counted_samples
            .checked_add(speaker.sample_count)
            .ok_or(ModelError::MalformedArtifact)?;

        let occupancy_sum = speaker.occupancy.iter().sum::<f64>();
        let sample_count = speaker.sample_count as f64;
        let tolerance = 1e-8 * sample_count.max(1.0);
        if !occupancy_sum.is_finite() || (occupancy_sum - sample_count).abs() > tolerance {
            return Err(ModelError::MalformedArtifact);
        }

        let expected = map_means(
            ubm_means,
            &speaker.occupancy,
            &speaker.first_moments,
            model.config.relevance,
        )
        .map_err(|_| ModelError::MalformedArtifact)?;

        for (((occupancy, first), actual), expected_component) in speaker
            .occupancy
            .iter()
            .zip(&speaker.first_moments)
            .zip(&speaker.adapted_means)
            .zip(&expected)
        {
            if !occupancy.is_finite()
                || *occupancy < 0.0
                || *occupancy > sample_count + tolerance
                || first.len() != dimension
                || actual.len() != dimension
                || expected_component.len() != dimension
                || first.iter().any(|value| !value.is_finite())
                || actual
                    .iter()
                    .zip(expected_component)
                    .any(|(value, expected_value)| {
                        !value.is_finite() || value.to_bits() != expected_value.to_bits()
                    })
            {
                return Err(ModelError::MalformedArtifact);
            }
        }

        let adapted = with_means(&model.ubm, speaker.adapted_means.clone())
            .map_err(|_| ModelError::MalformedArtifact)?;
        if log_likelihood(&adapted, &zero_row).is_err() {
            return Err(ModelError::MalformedArtifact);
        }
    }

    if counted_samples != model.training_sample_count {
        return Err(ModelError::MalformedArtifact);
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use kcode_speaker_types::ObjectId;
    use serde_json::{Value, json};

    fn key(value: &str) -> Key {
        Key::parse(value).unwrap()
    }

    fn vector(seed: u8) -> FeatureVector {
        let mut values = [0_u8; FEATURE_COUNT];
        for (index, value) in values.iter_mut().enumerate() {
            *value = u8::try_from((usize::from(seed) + index * 3) % 90).unwrap();
        }
        FeatureVector::new(values).unwrap()
    }

    fn sample(id: &str, cohort: &str, speaker: &str, seed: u8) -> LabeledSample {
        LabeledSample {
            sample_id: key(id),
            attempt_id: key(&format!("attempt:{id}")),
            speaker_id: key(speaker),
            cohort_id: key(cohort),
            group_id: key("group:fixture"),
            clip_object: ObjectId::parse(&format!("C{seed:07}")).unwrap(),
            primary_language: key("eng"),
            recording_kind: RecordingKind::VoiceNote,
            usable_speech_ms: 1_000 + u32::from(seed),
            recording_quality: 80 + seed % 20,
            features: vector(seed),
        }
    }

    fn config(mask: FeatureMask) -> ModelConfig {
        ModelConfig {
            mask,
            components: 1,
            relevance: 2.0,
            variance_floor: 0.01,
            absolute_threshold: -1e9,
            margin_threshold: -1e9,
        }
    }

    fn training() -> Vec<LabeledSample> {
        vec![
            sample("sample:1", "cohort", "alice", 10),
            sample("sample:2", "cohort", "alice", 14),
            sample("sample:3", "cohort", "bob", 60),
            sample("sample:4", "cohort", "bob", 66),
        ]
    }

    fn assert_malformed(value: &Value) {
        let bytes = serde_json::to_vec(value).unwrap();
        assert_eq!(decode(&bytes).err(), Some(ModelError::MalformedArtifact));
    }

    #[test]
    fn all_24_covers_exactly_the_shared_feature_contract() {
        assert_eq!(FEATURE_COUNT, 24);
        assert_eq!(u64::from(ALL_24), (1_u64 << FEATURE_COUNT) - 1);
    }

    #[test]
    fn minimal_rows_fit_without_media_or_attempt_metadata() {
        let cohort = key("cohort");
        let rows = training()
            .into_iter()
            .map(|sample| TrainingRow {
                sample_id: sample.sample_id,
                speaker_id: sample.speaker_id,
                cohort_id: sample.cohort_id,
                features: sample.features,
            })
            .collect::<Vec<_>>();
        let model = fit_rows(FitRowsInput {
            cohort_id: &cohort,
            rows: &rows,
            config: config(ALL_24),
        })
        .unwrap();
        let result = identify(&model, &cohort, &vector(12)).unwrap();
        assert_eq!(result.best.speaker_id, key("alice"));
    }

    #[test]
    fn hand_computed_map_and_unoccupied_component() {
        let ubm = vec![vec![1.0, -1.0], vec![7.0, 8.0]];
        let occupancy = vec![2.0, 0.0];
        let first_moments = vec![vec![6.0, 2.0], vec![0.0, 0.0]];
        let adapted = map_means(&ubm, &occupancy, &first_moments, 2.0).unwrap();

        assert_eq!(adapted[0], vec![2.0, 0.0]);
        assert_eq!(adapted[1], ubm[1]);
    }

    #[test]
    fn k1_llr_is_shared_covariance_mahalanobis_difference() {
        let samples = training();
        let cohort = key("cohort");
        let mask = FeatureMask::from_bits(1).unwrap();
        let model = fit(FitInput {
            cohort_id: &cohort,
            samples: &samples,
            config: config(mask),
        })
        .unwrap();

        let probe = vector(12);
        let result = identify(&model, &cohort, &probe).unwrap();
        let speaker = model
            .speakers
            .iter()
            .find(|speaker| speaker.speaker_id == result.best.speaker_id)
            .unwrap();
        let x = (f64::from(probe.as_ref()[0]) - model.normalizer_mean[0]) / model.normalizer_std[0];
        let ubm_mean = means(&model.ubm)[0][0];
        let value = serde_json::to_value(&model.ubm).unwrap();
        let variance = value["variances"][0][0].as_f64().unwrap();
        let adapted_mean = speaker.adapted_means[0][0];
        let expected =
            -0.5 * ((x - adapted_mean).powi(2) / variance - (x - ubm_mean).powi(2) / variance);

        assert!((result.best.llr - expected).abs() < 1e-12);
    }

    #[test]
    fn threshold_equality_and_one_speaker_rules() {
        let samples = training();
        let cohort = key("cohort");
        let mut model = fit(FitInput {
            cohort_id: &cohort,
            samples: &samples,
            config: config(ALL_24),
        })
        .unwrap();
        let probe = vector(12);
        let initial = identify(&model, &cohort, &probe).unwrap();
        let margin = initial.best.llr - initial.runner_up.as_ref().unwrap().llr;
        model.config.absolute_threshold = initial.best.llr;
        model.config.margin_threshold = margin;

        let equality = identify(&model, &cohort, &probe).unwrap();
        assert!(equality.absolute_pass);
        assert!(equality.margin_pass);
        assert!(matches!(equality.decision, Decision::Known { .. }));

        let one = fit(FitInput {
            cohort_id: &cohort,
            samples: &samples[..2],
            config: config(ALL_24),
        })
        .unwrap();
        let result = identify(&one, &cohort, &probe).unwrap();
        assert!(result.runner_up.is_none());
        assert!(result.margin_pass);

        let mut blocked = one;
        blocked.config.absolute_threshold = result.best.llr + 1e-12;
        assert!(matches!(
            identify(&blocked, &cohort, &probe).unwrap().decision,
            Decision::Unknown
        ));
    }

    #[test]
    fn cohort_mismatch_is_typed() {
        let samples = training();
        let cohort = key("cohort");
        let model = fit(FitInput {
            cohort_id: &cohort,
            samples: &samples,
            config: config(ALL_24),
        })
        .unwrap();

        assert_eq!(
            identify(&model, &key("different"), &vector(12)),
            Err(ModelError::CohortMismatch)
        );
    }

    #[test]
    fn input_order_does_not_change_the_artifact() {
        let samples = training();
        let cohort = key("cohort");
        let first = fit(FitInput {
            cohort_id: &cohort,
            samples: &samples,
            config: config(ALL_24),
        })
        .unwrap();

        let mut reversed = training();
        reversed.reverse();
        let second = fit(FitInput {
            cohort_id: &cohort,
            samples: &reversed,
            config: config(ALL_24),
        })
        .unwrap();

        assert_eq!(first.manifest_sha256, second.manifest_sha256);
        assert_eq!(encode(&first).unwrap(), encode(&second).unwrap());
    }

    #[test]
    fn same_sample_ids_with_changed_features_or_labels_change_identity() {
        let samples = training();
        let cohort = key("cohort");
        let original = fit(FitInput {
            cohort_id: &cohort,
            samples: &samples,
            config: config(ALL_24),
        })
        .unwrap();
        let original_bytes = encode(&original).unwrap();

        let mut feature_changed = training();
        let mut values = *feature_changed[0].features.as_ref();
        values[0] += 1;
        feature_changed[0].features = FeatureVector::new(values).unwrap();
        assert!(
            samples
                .iter()
                .zip(&feature_changed)
                .all(|(left, right)| left.sample_id == right.sample_id)
        );
        let feature_model = fit(FitInput {
            cohort_id: &cohort,
            samples: &feature_changed,
            config: config(ALL_24),
        })
        .unwrap();
        assert_ne!(original.manifest_sha256, feature_model.manifest_sha256);
        assert_ne!(original_bytes, encode(&feature_model).unwrap());

        let mut label_changed = training();
        label_changed[0].speaker_id = key("carol");
        assert!(
            samples
                .iter()
                .zip(&label_changed)
                .all(|(left, right)| left.sample_id == right.sample_id)
        );
        let label_model = fit(FitInput {
            cohort_id: &cohort,
            samples: &label_changed,
            config: config(ALL_24),
        })
        .unwrap();
        assert_ne!(original.manifest_sha256, label_model.manifest_sha256);
        assert_ne!(original_bytes, encode(&label_model).unwrap());
    }

    #[test]
    fn serialization_round_trip_is_deterministic_and_valid() {
        let samples = training();
        let cohort = key("cohort");
        let model = fit(FitInput {
            cohort_id: &cohort,
            samples: &samples,
            config: config(ALL_24),
        })
        .unwrap();

        let first = encode(&model).unwrap();
        let decoded = decode(&first).unwrap();
        let second = encode(&decoded).unwrap();
        let decoded_again = decode(&second).unwrap();
        let third = encode(&decoded_again).unwrap();

        assert_eq!(first, second);
        assert_eq!(second, third);
        assert_eq!(
            identify(&model, &cohort, &vector(12)).unwrap(),
            identify(&decoded_again, &cohort, &vector(12)).unwrap()
        );
    }

    #[test]
    fn decode_rejects_incompatible_and_corrupted_artifacts() {
        let samples = training();
        let cohort = key("cohort");
        let model = fit(FitInput {
            cohort_id: &cohort,
            samples: &samples,
            config: config(ALL_24),
        })
        .unwrap();
        let bytes = encode(&model).unwrap();
        let value = serde_json::from_slice::<Value>(&bytes).unwrap();

        assert_eq!(decode(b"{}").err(), Some(ModelError::MalformedArtifact));
        assert_eq!(decode(b"{").err(), Some(ModelError::MalformedArtifact));

        let mut old_version = value.clone();
        old_version["version"] = json!(1);
        assert_malformed(&old_version);

        let mut wrong_feature_count = value.clone();
        wrong_feature_count["feature_count"] = json!(35);
        assert_malformed(&wrong_feature_count);

        let mut bad_scale = value.clone();
        bad_scale["normalizer_std"][0] = json!(0.0);
        assert_malformed(&bad_scale);

        let mut bad_map_state = value;
        let adapted = bad_map_state["speakers"][0]["adapted_means"][0][0]
            .as_f64()
            .unwrap();
        bad_map_state["speakers"][0]["adapted_means"][0][0] = json!(adapted + 0.25);
        assert_malformed(&bad_map_state);
    }

    #[test]
    fn rejects_bad_samples_and_configuration() {
        let cohort = key("cohort");
        let empty = Vec::<LabeledSample>::new();
        assert_eq!(
            fit(FitInput {
                cohort_id: &cohort,
                samples: &empty,
                config: config(ALL_24),
            })
            .err(),
            Some(ModelError::InvalidSamples)
        );

        let mut mixed = training();
        mixed[0].cohort_id = key("other");
        assert_eq!(
            fit(FitInput {
                cohort_id: &cohort,
                samples: &mixed,
                config: config(ALL_24),
            })
            .err(),
            Some(ModelError::CohortMismatch)
        );

        let mut duplicate = training();
        duplicate[1].sample_id = duplicate[0].sample_id.clone();
        assert_eq!(
            fit(FitInput {
                cohort_id: &cohort,
                samples: &duplicate,
                config: config(ALL_24),
            })
            .err(),
            Some(ModelError::DuplicateSampleId)
        );

        let same = vec![
            sample("same:1", "cohort", "alice", 10),
            sample("same:2", "cohort", "alice", 10),
        ];
        assert_eq!(
            fit(FitInput {
                cohort_id: &cohort,
                samples: &same,
                config: config(ALL_24),
            })
            .err(),
            Some(ModelError::ZeroVariance)
        );

        let samples = training();
        for bad in [
            ModelConfig {
                components: 0,
                ..config(ALL_24)
            },
            ModelConfig {
                components: 5,
                ..config(ALL_24)
            },
            ModelConfig {
                relevance: 0.0,
                ..config(ALL_24)
            },
            ModelConfig {
                variance_floor: f64::NAN,
                ..config(ALL_24)
            },
            ModelConfig {
                absolute_threshold: f64::INFINITY,
                ..config(ALL_24)
            },
            ModelConfig {
                margin_threshold: f64::NAN,
                ..config(ALL_24)
            },
        ] {
            assert_eq!(
                fit(FitInput {
                    cohort_id: &cohort,
                    samples: &samples,
                    config: bad,
                })
                .err(),
                Some(ModelError::InvalidConfig)
            );
        }
    }

    #[test]
    fn rejects_nonconverged_gmm() {
        let rows = vec![
            vec![-5.0],
            vec![-4.0],
            vec![-1.0],
            vec![2.0],
            vec![3.0],
            vec![10.0],
        ];
        assert_eq!(
            fit_ubm_with_limits(&rows, 2, 0.01, 1, f64::MIN_POSITIVE,).err(),
            Some(ModelError::Nonconverged)
        );
    }
}