use anyhow::{Result, anyhow};
use qwen_asr::context::QwenModel;
use serde::{Deserialize, Serialize};
use crate::vector::cosine_similarity;
pub(crate) const WAKE_WORD_EMBEDDING_DIM: usize = 1024;
pub(crate) const WINDOW_MEL_FRAMES: usize = 76;
pub(crate) const WINDOW_SAMPLES: usize = WINDOW_MEL_FRAMES * 160;
const MAX_TOKENS_PER_CHUNK: usize = 13;
pub(crate) const SCORE_STRIDE_MEL_FRAMES: usize = 16;
const MAX_SOFT_FLOOR: f32 = 0.75;
const DEFAULT_SOFT_FLOOR: f32 = 0.55;
pub(crate) const ENROLLMENT_CONSISTENCY_MIN_SIMILARITY: f32 = 0.70;
pub(crate) const ENROLLMENT_CONSISTENCY_MIN_FRACTION: f32 = 0.7;
pub(crate) const MIN_ENROLLMENT_UTTERANCES: usize = 5;
pub(crate) fn encode_window(model: &QwenModel, samples: &[f32]) -> Result<Vec<f32>> {
let take = samples.len().min(WINDOW_SAMPLES);
let mut window = vec![0.0; WINDOW_SAMPLES];
window[WINDOW_SAMPLES - take..].copy_from_slice(&samples[samples.len() - take..]);
let (mel, mel_frames) = qwen_asr::audio::mel_spectrogram(&window).ok_or_else(|| {
anyhow!("mel_spectrogram returned None for {WINDOW_SAMPLES}-sample window")
})?;
debug_assert_eq!(mel_frames, WINDOW_MEL_FRAMES);
let (features, total_tokens) = model
.encoder
.forward(&model.config, &mel, mel_frames, None)
.ok_or_else(|| anyhow!("encoder forward failed for {mel_frames}-frame window"))?;
debug_assert!(total_tokens <= MAX_TOKENS_PER_CHUNK);
let output_dim = model.config.enc_output_dim;
let mut pooled = vec![0.0f32; output_dim];
for t in 0..total_tokens {
let base = t * output_dim;
for (d, v) in pooled.iter_mut().enumerate() {
*v += features[base + d];
}
}
#[expect(clippy::cast_precision_loss)]
let inv = 1.0 / total_tokens as f32;
for v in &mut pooled {
*v *= inv;
}
l2_normalize_in_place(&mut pooled);
Ok(pooled)
}
#[cfg(test)]
#[must_use]
pub(crate) fn l2_normalize(v: &[f32]) -> Vec<f32> {
let mut out = v.to_vec();
l2_normalize_in_place(&mut out);
out
}
fn l2_normalize_in_place(v: &mut [f32]) {
let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
if norm > 1e-8 {
let inv = 1.0 / norm;
for x in v.iter_mut() {
*x *= inv;
}
}
}
pub(crate) const ENROLLMENT_SCHEMA_VERSION: u32 = 2;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub(crate) struct Calibration {
pub(crate) neg_mean: f32,
pub(crate) neg_std: f32,
pub(crate) neg_p99: f32,
pub(crate) n_negatives: usize,
}
impl Calibration {
#[must_use]
pub(crate) fn soft_floor(&self) -> f32 {
self.neg_p99.clamp(DEFAULT_SOFT_FLOOR, MAX_SOFT_FLOOR)
}
#[must_use]
pub(crate) fn soft_score(&self, cosine: f32) -> f32 {
let floor = self.soft_floor();
((cosine - floor) / (1.0 - floor)).clamp(0.0, 1.0)
}
}
impl Default for Calibration {
fn default() -> Self {
Self {
neg_mean: 0.0,
neg_std: 0.0,
neg_p99: DEFAULT_SOFT_FLOOR,
n_negatives: 0,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub(crate) struct WakeWordEnrollment {
#[serde(default)]
pub(crate) schema_version: u32,
pub(crate) phrase: String,
pub(crate) embedding_dim: usize,
pub(crate) prototype: Vec<f32>,
pub(crate) utterance_count: usize,
#[serde(default)]
pub(crate) calibration: Calibration,
#[serde(default)]
pub(crate) negative_prototypes: Vec<Vec<f32>>,
#[serde(default)]
pub(crate) created_at: String,
#[serde(default)]
pub(crate) trained_at: String,
}
impl WakeWordEnrollment {
#[must_use]
pub(crate) fn build(
phrase: String,
utterance_embeddings: &[Vec<f32>],
calibration: Calibration,
negative_embeddings: &[Vec<f32>],
created_at: String,
trained_at: String,
) -> Option<Self> {
if utterance_embeddings.len() < MIN_ENROLLMENT_UTTERANCES {
return None;
}
let dim = WAKE_WORD_EMBEDDING_DIM;
let mut prototype = vec![0.0f32; dim];
for emb in utterance_embeddings {
if emb.len() != dim {
return None;
}
for (p, e) in prototype.iter_mut().zip(emb) {
*p += e;
}
}
#[expect(clippy::cast_precision_loss)]
let inv = 1.0 / utterance_embeddings.len() as f32;
for p in &mut prototype {
*p *= inv;
}
l2_normalize_in_place(&mut prototype);
Some(Self {
schema_version: ENROLLMENT_SCHEMA_VERSION,
phrase,
embedding_dim: dim,
prototype,
utterance_count: utterance_embeddings.len(),
calibration,
negative_prototypes: distill_negative_prototypes(negative_embeddings),
created_at,
trained_at,
})
}
#[must_use]
pub(crate) fn cosine(&self, embedding: &[f32]) -> f32 {
cosine_similarity(embedding, &self.prototype)
}
#[must_use]
pub(crate) fn max_negative_cosine(&self, embedding: &[f32]) -> f32 {
self.negative_prototypes
.iter()
.map(|a| cosine_similarity(embedding, a))
.fold(0.0_f32, f32::max)
}
#[must_use]
pub(crate) fn soft_score(&self, embedding: &[f32]) -> f32 {
let pos = self.cosine(embedding);
let neg = self.max_negative_cosine(embedding);
if !self.negative_prototypes.is_empty() && neg > pos {
return 0.0;
}
self.calibration.soft_score(pos)
}
}
pub(crate) const MAX_NEGATIVE_PROTOTYPES: usize = 8;
#[must_use]
pub(crate) fn distill_negative_prototypes(negatives: &[Vec<f32>]) -> Vec<Vec<f32>> {
if negatives.is_empty() {
return Vec::new();
}
if negatives.len() == 1 {
return vec![negatives[0].clone()];
}
let k = negatives.len().min(MAX_NEGATIVE_PROTOTYPES);
let mut chosen: Vec<usize> = Vec::with_capacity(k);
chosen.push(0);
while chosen.len() < k {
let mut best_idx = None;
let mut best_sim = f32::MAX;
for (i, n) in negatives.iter().enumerate() {
if chosen.contains(&i) {
continue;
}
let min_sim = chosen
.iter()
.map(|&c| cosine_similarity(n, &negatives[c]))
.fold(f32::MAX, f32::min);
if min_sim < best_sim {
best_sim = min_sim;
best_idx = Some(i);
}
}
match best_idx {
Some(i) => chosen.push(i),
None => break,
}
}
chosen.into_iter().map(|i| negatives[i].clone()).collect()
}
#[must_use]
#[expect(
clippy::cast_possible_truncation,
clippy::cast_precision_loss,
clippy::cast_sign_loss
)]
pub(crate) fn calibrate_negatives(prototype: &[f32], negatives: &[Vec<f32>]) -> Calibration {
if negatives.is_empty() {
return Calibration::default();
}
let mut cosines: Vec<f32> = negatives
.iter()
.map(|n| cosine_similarity(n, prototype))
.collect();
cosines.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
let n = cosines.len();
let mean = cosines.iter().sum::<f32>() / n as f32;
let var = cosines.iter().map(|c| (c - mean) * (c - mean)).sum::<f32>() / n as f32;
let std = var.sqrt();
let p99 = if n >= 4 {
let idx = ((n - 1) as f32 * 0.99).round() as usize;
cosines[idx.min(n - 1)]
} else {
(mean + 2.0 * std).min(cosines[n - 1])
};
Calibration {
neg_mean: mean,
neg_std: std,
neg_p99: p99.clamp(DEFAULT_SOFT_FLOOR, MAX_SOFT_FLOOR),
n_negatives: n,
}
}
#[cfg(test)]
mod tests {
use super::*;
fn norm(v: &[f32]) -> Vec<f32> {
l2_normalize(v)
}
fn unit_with_cosine(reference: &[f32], c: f32, seed: u64) -> Vec<f32> {
use rand::{RngExt, SeedableRng};
let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
let mut r: Vec<f32> = (0..reference.len())
.map(|_| rng.random::<f32>() * 2.0 - 1.0)
.collect();
let proj: f32 = reference.iter().zip(&r).map(|(a, b)| a * b).sum();
for (ri, &ri_ref) in r.iter_mut().zip(reference) {
*ri -= proj * ri_ref;
}
let ortho = l2_normalize(&r);
let s = (1.0 - c * c).max(0.0).sqrt();
let mut v: Vec<f32> = reference
.iter()
.zip(&ortho)
.map(|(&a, &b)| c * a + s * b)
.collect();
l2_normalize_in_place(&mut v);
v
}
#[test]
fn calibration_floor_maps_negatives_to_zero() {
#[expect(clippy::cast_precision_loss)] let proto = norm(
&(0..WAKE_WORD_EMBEDDING_DIM)
.map(|i| i as f32)
.collect::<Vec<_>>(),
);
let negatives: Vec<Vec<f32>> = (0..20)
.map(|i| {
#[expect(clippy::cast_precision_loss)] unit_with_cosine(
&proto,
0.20 + i as f32 * 0.0125,
1000 + u64::try_from(i).unwrap(),
)
})
.collect();
for (i, n) in negatives.iter().enumerate() {
let c = cosine_similarity(n, &proto);
#[expect(clippy::cast_precision_loss)] let expected = 0.20 + i as f32 * 0.0125;
assert!(
(c - expected).abs() < 1e-2,
"negative {i} cosine {c} != expected {expected}"
);
}
let cal = calibrate_negatives(&proto, &negatives);
assert!(
(cal.neg_p99 - DEFAULT_SOFT_FLOOR).abs() < 1e-6,
"p99 below the floor must clamp to the default floor, got {}",
cal.neg_p99,
);
#[expect(clippy::float_cmp)] {
assert_eq!(cal.soft_floor(), DEFAULT_SOFT_FLOOR);
}
for n in &negatives {
#[expect(clippy::float_cmp)] {
assert_eq!(cal.soft_score(cosine_similarity(n, &proto)), 0.0);
}
}
let match_emb = unit_with_cosine(&proto, 0.9, 999);
assert!(cal.soft_score(cosine_similarity(&match_emb, &proto)) > 0.5);
let hot_negatives: Vec<Vec<f32>> = (0..10)
.map(|i| {
#[expect(clippy::cast_precision_loss)] unit_with_cosine(
&proto,
0.69 + i as f32 * 0.002,
2000 + u64::try_from(i).unwrap(),
)
})
.collect();
let hot = calibrate_negatives(&proto, &hot_negatives);
assert!(
hot.soft_floor() > DEFAULT_SOFT_FLOOR + 0.1,
"floor must track the high negative band, got {}",
hot.soft_floor(),
);
for n in &hot_negatives {
#[expect(clippy::float_cmp)] {
assert_eq!(hot.soft_score(cosine_similarity(n, &proto)), 0.0);
}
}
let strong = unit_with_cosine(&proto, 0.9, 998);
assert!(hot.soft_score(cosine_similarity(&strong, &proto)) > 0.5);
}
#[test]
fn anti_prototype_gate_rejects_windows_closer_to_a_negative() {
#[expect(clippy::cast_precision_loss)] let proto = norm(
&(0..WAKE_WORD_EMBEDDING_DIM)
.map(|i| i as f32)
.collect::<Vec<_>>(),
);
let confusable = unit_with_cosine(&proto, 0.85, 42);
let cal = Calibration::default(); assert!(cal.soft_score(cosine_similarity(&confusable, &proto)) > 0.5);
let five: Vec<Vec<f32>> = vec![proto.clone(); MIN_ENROLLMENT_UTTERANCES];
let enr = WakeWordEnrollment::build(
"mahbot".into(),
&five,
cal.clone(),
std::slice::from_ref(&confusable),
String::new(),
String::new(),
)
.expect("enrollment");
#[expect(clippy::float_cmp)] {
assert_eq!(enr.soft_score(&confusable), 0.0);
}
assert!(enr.soft_score(&proto) > 0.5);
let no_neg = WakeWordEnrollment::build(
"mahbot".into(),
&five,
cal.clone(),
&[],
String::new(),
String::new(),
)
.expect("enrollment");
assert!(no_neg.soft_score(&confusable) > 0.5);
}
#[test]
fn distill_negative_prototypes_spreads_selection() {
#[expect(clippy::cast_precision_loss)] let proto = norm(
&(0..WAKE_WORD_EMBEDDING_DIM)
.map(|i| i as f32)
.collect::<Vec<_>>(),
);
let negatives: Vec<Vec<f32>> = (0..24)
.map(|i| {
#[expect(clippy::cast_precision_loss)] unit_with_cosine(
&proto,
0.30 + i as f32 * 0.01,
3000 + u64::try_from(i).unwrap(),
)
})
.collect();
let distilled = distill_negative_prototypes(&negatives);
assert!(!distilled.is_empty());
assert!(distilled.len() <= MAX_NEGATIVE_PROTOTYPES);
assert!(
distilled.len() >= 2,
"24 spread negatives should yield >1 prototype, got {}",
distilled.len(),
);
let mut max_pairwise = 0.0f32;
for (i, a) in distilled.iter().enumerate() {
for b in distilled.iter().skip(i + 1) {
max_pairwise = max_pairwise.max(cosine_similarity(a, b));
}
}
assert!(
max_pairwise < 0.9,
"distilled prototypes should be spread, max pairwise cosine {max_pairwise}",
);
for d in &distilled {
assert!(
negatives
.iter()
.any(|n| (cosine_similarity(n, d) - 1.0).abs() < 1e-4)
);
}
let single = distill_negative_prototypes(&negatives[..1]);
assert_eq!(single.len(), 1);
assert!((cosine_similarity(&single[0], &negatives[0]) - 1.0).abs() < 1e-4);
assert!(distill_negative_prototypes(&[]).is_empty());
}
#[test]
fn calibration_default_floor() {
let cal = Calibration::default();
#[expect(clippy::float_cmp)] {
assert_eq!(cal.soft_floor(), DEFAULT_SOFT_FLOOR);
}
#[expect(clippy::float_cmp)] {
assert_eq!(cal.soft_score(DEFAULT_SOFT_FLOOR), 0.0);
}
#[expect(clippy::float_cmp)] {
assert_eq!(cal.soft_score(1.0), 1.0);
}
}
#[test]
fn build_requires_min_utterances() {
let emb = norm(&vec![1.0; WAKE_WORD_EMBEDDING_DIM]);
let enough: Vec<Vec<f32>> = vec![emb.clone(); MIN_ENROLLMENT_UTTERANCES];
assert!(
WakeWordEnrollment::build(
"mahbot".into(),
&enough,
Calibration::default(),
&[],
String::new(),
String::new(),
)
.is_some()
);
let too_few: Vec<Vec<f32>> = vec![emb; MIN_ENROLLMENT_UTTERANCES - 1];
assert!(
WakeWordEnrollment::build(
"mahbot".into(),
&too_few,
Calibration::default(),
&[],
String::new(),
String::new(),
)
.is_none()
);
}
#[test]
fn prototype_is_unit_norm() {
#[expect(clippy::cast_precision_loss)] let emb = norm(
&(0..WAKE_WORD_EMBEDDING_DIM)
.map(|i| i as f32)
.collect::<Vec<_>>(),
);
let five: Vec<Vec<f32>> = vec![emb; MIN_ENROLLMENT_UTTERANCES];
let enr = WakeWordEnrollment::build(
"mahbot".into(),
&five,
Calibration::default(),
&[],
String::new(),
String::new(),
)
.expect("enrollment");
let norm: f32 = enr.prototype.iter().map(|x| x * x).sum::<f32>().sqrt();
assert!((norm - 1.0).abs() < 1e-3);
assert_eq!(enr.schema_version, ENROLLMENT_SCHEMA_VERSION);
}
#[test]
fn serde_roundtrip_v2() {
#[expect(clippy::cast_precision_loss)] let emb = norm(
&(0..WAKE_WORD_EMBEDDING_DIM)
.map(|i| i as f32)
.collect::<Vec<_>>(),
);
let eight: Vec<Vec<f32>> = vec![emb; 8];
let enr = WakeWordEnrollment::build(
"mahbot".into(),
&eight,
Calibration::default(),
&[],
"created".into(),
"trained".into(),
)
.expect("enrollment");
let json = serde_json::to_string(&enr).expect("serialize");
let back: WakeWordEnrollment = serde_json::from_str(&json).expect("deserialize");
assert_eq!(back.schema_version, ENROLLMENT_SCHEMA_VERSION);
assert_eq!(back.embedding_dim, WAKE_WORD_EMBEDDING_DIM);
assert_eq!(back.prototype.len(), WAKE_WORD_EMBEDDING_DIM);
assert_eq!(back.phrase, "mahbot");
}
#[test]
fn v1_schema_rejected() {
let v1 = r#"{"schema_version":1,"phrase":"mahbot","embedding_dim":96,
"window_size":3,"classifier":[]}"#;
let parsed: serde_json::Result<WakeWordEnrollment> = serde_json::from_str(v1);
assert!(parsed.is_err());
}
}