#![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,
}
#[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 selected = mask_indices(input.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,
input.config.components,
input.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(input.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,
input.config.relevance,
)?;
}
let manifest_sha256 = manifest(&samples);
let model = ModelSnapshot {
version: SNAPSHOT_VERSION,
feature_count: WIRE_FEATURE_COUNT,
cohort_id: input.cohort_id.clone(),
config: input.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 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 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)
);
}
}