Skip to main content

max / audiofiles

17.3 KB · 490 lines History Blame Raw
1 //! Optional trained head (Phase 6 of the hybrid classifier).
2 //!
3 //! For large libraries, brute-force k-NN ([`super::exemplar`]) costs O(N·d) per query.
4 //! This module *distills* the same exemplar store into a compact, per-library
5 //! **one-vs-rest logistic-regression** model: per tag, a [`NUM_FEATURES`]-weight vector +
6 //! bias scoring the standardized feature vector through a sigmoid. Inference is then
7 //! O(tags·d), independent of library size, and the whole model is a few KB.
8 //!
9 //! The head is a **local, regenerable cache** (the `trained_head` table, not synced):
10 //! it rebuilds from `sample_features` + `tags` on any device, so there is nothing to sync
11 //! and nothing data-derived to ship. It is *optional*, the k-NN layer stays the source of
12 //! truth and the only path offering neighbor explainability; the head trades that for speed
13 //! on big libraries. Training is fully deterministic (zero init, fixed-iteration full-batch
14 //! gradient descent, no RNG), so two devices distilling the same exemplars agree.
15
16 use std::collections::HashMap;
17
18 use crate::db::Database;
19 use crate::error::{CoreError, Result, unix_now};
20 use tracing::instrument;
21
22 use super::classify::{FEATURE_VERSION, NUM_FEATURES};
23 use super::exemplar::{
24 self, MlOutcome, TagScore, apply_policy, load_query_vector, sample_tag_set, standardize_vec,
25 };
26
27 /// A tag needs at least this many positive exemplars to earn a distilled model; rarer
28 /// tags stay with k-NN (too few points to fit a stable linear boundary).
29 pub const MIN_POSITIVE_EXAMPLES: usize = 5;
30
31 /// Below this many labeled exemplars, brute-force k-NN is already fast; the head is
32 /// recommended only once the library grows past it. Advisory, callers may train anyway.
33 pub const RECOMMENDED_MIN_EXEMPLARS: usize = 2_000;
34
35 // Training hyperparameters, deterministic full-batch gradient descent.
36 const ITERATIONS: usize = 300;
37 const LEARNING_RATE: f64 = 0.5;
38 const L2: f64 = 1e-3;
39
40 /// A per-tag linear scorer: `sigmoid(weights · x_std + bias)`.
41 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
42 pub struct TagModel {
43 pub tag: String,
44 pub weights: Vec<f64>,
45 pub bias: f64,
46 }
47
48 /// A compact, per-library classifier distilled from the exemplar store.
49 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
50 pub struct TrainedHead {
51 /// Feature version the model was trained against; stale versions are unusable.
52 pub feat_version: u32,
53 /// Standardization params (over the training exemplars), the head expects raw
54 /// feature vectors and standardizes them itself, matching k-NN.
55 pub means: Vec<f64>,
56 pub stds: Vec<f64>,
57 pub classes: Vec<TagModel>,
58 /// Number of labeled exemplars the model was trained on (staleness/worthwhile UI).
59 pub exemplar_count: usize,
60 pub trained_at: i64,
61 }
62
63 impl TrainedHead {
64 /// Whether the model matches the current feature layout. A mismatch means the vectors
65 /// were recomputed under a new `FEATURE_VERSION`; the head must be retrained.
66 pub fn is_current(&self) -> bool {
67 self.feat_version == FEATURE_VERSION
68 }
69
70 /// Score every modeled tag for a raw (un-standardized) query vector, highest first.
71 /// Neighbors are empty: the head trades k-NN's explainability for speed.
72 pub fn score(&self, raw_query: &[f64]) -> Vec<TagScore> {
73 if raw_query.len() != NUM_FEATURES {
74 return Vec::new();
75 }
76 let x = standardize_vec(raw_query, &self.means, &self.stds);
77 let mut scores: Vec<TagScore> = self
78 .classes
79 .iter()
80 .filter(|m| m.weights.len() == x.len())
81 .map(|m| {
82 let z = dot(&m.weights, &x) + m.bias;
83 TagScore {
84 tag: m.tag.clone(),
85 score: sigmoid(z),
86 neighbors: Vec::new(),
87 }
88 })
89 .collect();
90 // total_cmp never panics on NaN; the head is trained from the same
91 // finiteness-filtered exemplar set, so this is a backstop.
92 scores.sort_by(|a, b| b.score.total_cmp(&a.score).then(a.tag.cmp(&b.tag)));
93 scores
94 }
95 }
96
97 fn sigmoid(z: f64) -> f64 {
98 1.0 / (1.0 + (-z).exp())
99 }
100
101 fn dot(a: &[f64], b: &[f64]) -> f64 {
102 a.iter().zip(b).map(|(x, y)| x * y).sum()
103 }
104
105 /// Fit one balanced binary logistic regression by deterministic full-batch gradient
106 /// descent with L2 regularization. Per-sample class weights counter the heavy
107 /// positive/negative imbalance of multi-label tags (a tag is rare among all exemplars).
108 fn train_one(xs: &[&[f64]], labels: &[bool], dims: usize) -> (Vec<f64>, f64) {
109 let n = xs.len() as f64;
110 let pos = labels.iter().filter(|y| **y).count() as f64;
111 let neg = n - pos;
112 // Balance so positives and negatives contribute equal total weight.
113 let w_pos = if pos > 0.0 { n / (2.0 * pos) } else { 0.0 };
114 let w_neg = if neg > 0.0 { n / (2.0 * neg) } else { 0.0 };
115
116 let mut w = vec![0.0; dims];
117 let mut b = 0.0;
118 for _ in 0..ITERATIONS {
119 let mut grad_w = vec![0.0; dims];
120 let mut grad_b = 0.0;
121 for (x, &y) in xs.iter().zip(labels) {
122 let p = sigmoid(dot(&w, x) + b);
123 let sw = if y { w_pos } else { w_neg };
124 let err = sw * (p - if y { 1.0 } else { 0.0 });
125 for d in 0..dims {
126 grad_w[d] += err * x[d];
127 }
128 grad_b += err;
129 }
130 for d in 0..dims {
131 // Mean gradient (weights sum to n) + L2 pull toward zero.
132 w[d] -= LEARNING_RATE * (grad_w[d] / n + L2 * w[d]);
133 }
134 b -= LEARNING_RATE * (grad_b / n);
135 }
136 (w, b)
137 }
138
139 /// Distill the current exemplar store into a trained head **without persisting it**.
140 /// Returns `None` when there is nothing worth training (no labeled exemplars, or no tag
141 /// clears [`MIN_POSITIVE_EXAMPLES`]).
142 #[instrument(skip_all)]
143 pub fn train(db: &Database) -> Result<Option<TrainedHead>> {
144 let index = exemplar::build_index(db)?;
145 if index.is_empty() {
146 return Ok(None);
147 }
148
149 // Shared standardized design matrix + per-row tag sets (borrow the index once).
150 let xs: Vec<&[f64]> = index.rows().map(|(v, _)| v).collect();
151 let row_tags: Vec<&[String]> = index.rows().map(|(_, t)| t).collect();
152
153 // Tags with enough positives to model; sorted for deterministic class order.
154 let mut counts: HashMap<&str, usize> = HashMap::new();
155 for tags in &row_tags {
156 for t in *tags {
157 *counts.entry(t.as_str()).or_default() += 1;
158 }
159 }
160 let mut candidates: Vec<&str> = counts
161 .iter()
162 .filter(|(_, c)| **c >= MIN_POSITIVE_EXAMPLES)
163 .map(|(t, _)| *t)
164 .collect();
165 candidates.sort_unstable();
166 if candidates.is_empty() {
167 return Ok(None);
168 }
169
170 let classes = candidates
171 .iter()
172 .map(|&tag| {
173 let labels: Vec<bool> = row_tags
174 .iter()
175 .map(|tags| tags.iter().any(|t| t == tag))
176 .collect();
177 let (weights, bias) = train_one(&xs, &labels, NUM_FEATURES);
178 TagModel {
179 tag: tag.to_string(),
180 weights,
181 bias,
182 }
183 })
184 .collect();
185
186 Ok(Some(TrainedHead {
187 feat_version: FEATURE_VERSION,
188 means: index.means().to_vec(),
189 stds: index.stds().to_vec(),
190 classes,
191 exemplar_count: index.len(),
192 trained_at: unix_now(),
193 }))
194 }
195
196 /// Distill and persist the head (the background-worker entry point). Persists the fresh
197 /// model, or [`clear`]s the stored head when there is nothing to train. Returns what was
198 /// stored (or `None`).
199 #[instrument(skip_all)]
200 pub fn train_and_save(db: &Database) -> Result<Option<TrainedHead>> {
201 match train(db)? {
202 Some(head) => {
203 save(db, &head)?;
204 Ok(Some(head))
205 }
206 None => {
207 clear(db)?;
208 Ok(None)
209 }
210 }
211 }
212
213 /// Persist the singleton trained head, replacing any prior model.
214 #[instrument(skip_all)]
215 pub fn save(db: &Database, head: &TrainedHead) -> Result<()> {
216 let json = serde_json::to_string(head).map_err(|e| CoreError::Serialization(e.to_string()))?;
217 db.conn().execute(
218 "INSERT INTO trained_head (id, feat_version, exemplar_count, model, trained_at)
219 VALUES (1, ?1, ?2, ?3, ?4)
220 ON CONFLICT(id) DO UPDATE SET
221 feat_version = ?1, exemplar_count = ?2, model = ?3, trained_at = ?4",
222 rusqlite::params![
223 head.feat_version,
224 head.exemplar_count as i64,
225 json,
226 head.trained_at
227 ],
228 )?;
229 Ok(())
230 }
231
232 /// Load the stored head, if any. A model trained under a stale `FEATURE_VERSION` is
233 /// dropped (returns `None`), callers should retrain.
234 #[instrument(skip_all)]
235 pub fn load(db: &Database) -> Result<Option<TrainedHead>> {
236 let json: Option<String> = db
237 .conn()
238 .query_row("SELECT model FROM trained_head WHERE id = 1", [], |row| {
239 row.get(0)
240 })
241 .ok();
242 let Some(json) = json else { return Ok(None) };
243 let head: TrainedHead =
244 serde_json::from_str(&json).map_err(|e| CoreError::Serialization(e.to_string()))?;
245 if !head.is_current() {
246 return Ok(None);
247 }
248 Ok(Some(head))
249 }
250
251 /// Remove the stored head.
252 #[instrument(skip_all)]
253 pub fn clear(db: &Database) -> Result<()> {
254 db.conn()
255 .execute("DELETE FROM trained_head WHERE id = 1", [])?;
256 Ok(())
257 }
258
259 /// Count current-version labeled exemplars cheaply (drives the worthwhile/staleness
260 /// decision without building the full index).
261 #[instrument(skip_all)]
262 pub fn labeled_exemplar_count(db: &Database) -> Result<usize> {
263 let n: i64 = db.conn().query_row(
264 "SELECT COUNT(DISTINCT t.sample_hash) FROM tags t
265 JOIN sample_features f ON f.hash = t.sample_hash AND f.feat_version = ?1
266 WHERE NOT EXISTS (
267 SELECT 1 FROM tag_provenance p
268 WHERE p.sample_hash = t.sample_hash AND p.tag = t.tag AND p.source = 'ml')",
269 [FEATURE_VERSION],
270 |row| row.get(0),
271 )?;
272 Ok(n as usize)
273 }
274
275 /// Whether a trained head is worth building for the current library size.
276 #[instrument(skip_all)]
277 pub fn is_worthwhile(db: &Database) -> Result<bool> {
278 Ok(labeled_exemplar_count(db)? >= RECOMMENDED_MIN_EXEMPLARS)
279 }
280
281 /// Score a stored sample against the head under each tag's policy, writing nothing.
282 /// Mirrors [`exemplar::preview_sample`] but via the distilled model.
283 #[instrument(skip_all)]
284 pub fn preview_sample(db: &Database, hash: &str, head: &TrainedHead) -> Result<MlOutcome> {
285 let Some(raw) = load_query_vector(db, hash)? else {
286 return Ok(MlOutcome::default());
287 };
288 let scores = head.score(&raw);
289 let existing = sample_tag_set(db, hash)?;
290 apply_policy(db, scores, &existing)
291 }
292
293 /// Score a stored sample via the head and **auto-apply** above-threshold tags
294 /// (`source = 'ml'`), returning what was applied and what's pending review.
295 #[instrument(skip_all)]
296 pub fn apply_ml_suggestions(db: &Database, hash: &str, head: &TrainedHead) -> Result<MlOutcome> {
297 let outcome = preview_sample(db, hash, head)?;
298 db.transaction_core(|tx| {
299 for s in &outcome.auto_applied {
300 crate::rules::apply_tag_sourced(db, tx, hash, &s.tag, "ml")?;
301 }
302 Ok(())
303 })?;
304 Ok(outcome)
305 }
306
307 /// Auto-apply above-threshold head tags to every analyzed, non-deleted sample, returning
308 /// the total tags applied. The head-routed counterpart to [`exemplar::auto_apply_library`]
309 /// O(tags) per sample, independent of library size.
310 #[instrument(skip_all)]
311 pub fn auto_apply_library(db: &Database, head: &TrainedHead) -> Result<usize> {
312 let hashes: Vec<String> = {
313 let mut stmt = db.conn().prepare(
314 "SELECT s.hash FROM live_samples s
315 JOIN sample_features f ON f.hash = s.hash AND f.feat_version = ?1",
316 )?;
317 stmt.query_map([FEATURE_VERSION], |row| row.get(0))?
318 .collect::<std::result::Result<Vec<_>, _>>()?
319 };
320 let mut applied = 0;
321 for hash in hashes {
322 applied += apply_ml_suggestions(db, &hash, head)?.auto_applied.len();
323 }
324 Ok(applied)
325 }
326
327 #[cfg(test)]
328 mod tests {
329 use super::*;
330
331 fn insert(db: &Database, hash: &str, vec: &[f64; NUM_FEATURES], tags: &[&str]) {
332 db.conn()
333 .execute(
334 "INSERT INTO samples (hash, original_name, file_extension, file_size, import_date, last_modified) \
335 VALUES (?1, ?1, 'wav', 1, 0, 0)",
336 [hash],
337 )
338 .unwrap();
339 let json = serde_json::to_string(&vec.to_vec()).unwrap();
340 db.conn()
341 .execute(
342 "INSERT INTO sample_features (hash, feat_version, vector, computed_at) VALUES (?1, ?2, ?3, 0)",
343 rusqlite::params![hash, FEATURE_VERSION, json],
344 )
345 .unwrap();
346 for t in tags {
347 crate::tags::add_tag(db, hash, t).unwrap();
348 }
349 }
350
351 /// Build two separated clusters of kicks and snares with `n` each, varied per index so
352 /// the standardizer sees non-zero variance.
353 fn seed_two_classes(db: &Database, n: usize) {
354 for i in 0..n {
355 let k = 0.01 * i as f64;
356 insert(
357 db,
358 &format!("k{i}"),
359 &[k; NUM_FEATURES],
360 &["instrument.drum.kick"],
361 );
362 insert(
363 db,
364 &format!("s{i}"),
365 &[10.0 + k; NUM_FEATURES],
366 &["instrument.drum.snare"],
367 );
368 }
369 }
370
371 #[test]
372 fn no_exemplars_trains_nothing() {
373 let db = Database::open_in_memory().unwrap();
374 assert!(train(&db).unwrap().is_none());
375 }
376
377 #[test]
378 fn too_few_positives_trains_nothing() {
379 let db = Database::open_in_memory().unwrap();
380 // Only 2 of each tag (< MIN_POSITIVE_EXAMPLES).
381 seed_two_classes(&db, 2);
382 assert!(train(&db).unwrap().is_none());
383 }
384
385 #[test]
386 fn distinguishes_two_classes() {
387 let db = Database::open_in_memory().unwrap();
388 seed_two_classes(&db, 8);
389 let head = train(&db).unwrap().expect("should train");
390 assert_eq!(head.classes.len(), 2);
391 assert!(head.is_current());
392
393 // A query in the kick cluster scores kick above snare.
394 let near_kick = [0.04; NUM_FEATURES];
395 let scores = head.score(&near_kick);
396 let kick = scores
397 .iter()
398 .find(|s| s.tag == "instrument.drum.kick")
399 .unwrap();
400 let snare = scores
401 .iter()
402 .find(|s| s.tag == "instrument.drum.snare")
403 .unwrap();
404 assert!(
405 kick.score > snare.score,
406 "kick {} vs snare {}",
407 kick.score,
408 snare.score
409 );
410 assert!(kick.score > 0.5);
411 }
412
413 #[test]
414 fn save_load_round_trip_and_clear() {
415 let db = Database::open_in_memory().unwrap();
416 seed_two_classes(&db, 6);
417 let head = train_and_save(&db).unwrap().unwrap();
418 let loaded = load(&db).unwrap().expect("persisted");
419 assert_eq!(loaded.classes.len(), head.classes.len());
420 assert_eq!(loaded.exemplar_count, head.exemplar_count);
421 clear(&db).unwrap();
422 assert!(load(&db).unwrap().is_none());
423 }
424
425 #[test]
426 fn stale_feature_version_is_ignored_on_load() {
427 let db = Database::open_in_memory().unwrap();
428 seed_two_classes(&db, 6);
429 let mut head = train(&db).unwrap().unwrap();
430 head.feat_version = FEATURE_VERSION + 1;
431 save(&db, &head).unwrap();
432 // Stored row exists but is stale → load returns None.
433 assert!(load(&db).unwrap().is_none());
434 }
435
436 #[test]
437 fn deterministic_across_runs() {
438 let train_once = || {
439 let db = Database::open_in_memory().unwrap();
440 seed_two_classes(&db, 7);
441 train(&db).unwrap().unwrap().classes
442 };
443 let a = train_once();
444 let b = train_once();
445 assert_eq!(a.len(), b.len());
446 for (ca, cb) in a.iter().zip(&b) {
447 assert_eq!(ca.tag, cb.tag);
448 assert_eq!(
449 ca.bias.to_bits(),
450 cb.bias.to_bits(),
451 "bias differs for {}",
452 ca.tag
453 );
454 for (wa, wb) in ca.weights.iter().zip(&cb.weights) {
455 assert_eq!(wa.to_bits(), wb.to_bits());
456 }
457 }
458 }
459
460 #[test]
461 fn apply_auto_tags_via_head() {
462 let db = Database::open_in_memory().unwrap();
463 seed_two_classes(&db, 10);
464 let head = train(&db).unwrap().unwrap();
465 // Unlabeled query in the kick cluster.
466 insert(&db, "q", &[0.05; NUM_FEATURES], &[]);
467 let outcome = apply_ml_suggestions(&db, "q", &head).unwrap();
468 let applied: Vec<&str> = outcome
469 .auto_applied
470 .iter()
471 .map(|s| s.tag.as_str())
472 .collect();
473 assert!(applied.contains(&"instrument.drum.kick"), "got {applied:?}");
474 // Provenance is 'ml'.
475 let prov = crate::rules::sample_tag_provenance(&db, "q").unwrap();
476 assert!(
477 prov.iter()
478 .any(|(t, s, _)| t == "instrument.drum.kick" && s == "ml")
479 );
480 }
481
482 #[test]
483 fn labeled_count_and_worthwhile() {
484 let db = Database::open_in_memory().unwrap();
485 seed_two_classes(&db, 5); // 10 labeled exemplars
486 assert_eq!(labeled_exemplar_count(&db).unwrap(), 10);
487 assert!(!is_worthwhile(&db).unwrap()); // well under RECOMMENDED_MIN_EXEMPLARS
488 }
489 }
490