Skip to main content

max / audiofiles

Add optional trained head + auto-tag routing (classifier Phase 6) Distil the tagged library into a compact per-library one-vs-rest logistic regression (analysis/trained_head.rs) as a faster alternative to brute-force k-NN on large libraries. Deterministic full-batch GD, class-balanced for the rare-positive multi-label case. Stored in a local, non-synced trained_head table (migration 024) as a regenerable cache. The heavy library-wide auto-tagging pass routes through the head when one is current (O(tags)/sample); detail-panel suggestions stay on k-NN for neighbor explainability. Settings > Auto-Tagging gains a "Trained model (optional)" subsection with Train/Retrain/Clear and a size-aware hint. Phase 7's classifier_layers shifts to migration 025. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Author: Max Johnson <me@maxj.phd> · 2026-06-15 01:03 UTC
Commit: 074e6d824ac53310a0c3df570b00b5735930c807
Parent: ebfd5ea
10 files changed, +732 insertions, -12 deletions
M crates/audiofiles-browser/src/backend/direct.rs +40 -2
@@ -26,7 +26,7 @@ use parking_lot::Mutex;
26 26
27 27 use super::{
28 28 Backend, BackendError, BackendEvent, BackendResult, ExportConfigDesc, ExportItemDesc,
29 - ImportStrategyDesc, ImportedFolderDesc,
29 + ImportStrategyDesc, ImportedFolderDesc, TrainedHeadInfo,
30 30 };
31 31
32 32 use crate::cleanup::{CleanupCommand, CleanupHandle};
@@ -639,8 +639,46 @@ impl Backend for DirectBackend {
639 639 }
640 640
641 641 fn ml_auto_apply_all(&self, k: usize) -> BackendResult<usize> {
642 + use audiofiles_core::analysis::{exemplar, trained_head};
643 + let db = self.db.lock();
644 + // Route the heavy library-wide pass through the distilled head when one is
645 + // current (O(tags) per sample instead of O(library)); else brute-force k-NN.
646 + // Detail-panel suggestions stay on k-NN for neighbor explainability.
647 + if let Some(head) = trained_head::load(&db)? {
648 + Ok(trained_head::auto_apply_library(&db, &head)?)
649 + } else {
650 + Ok(exemplar::auto_apply_library(&db, k)?)
651 + }
652 + }
653 +
654 + fn train_classifier_head(&self) -> BackendResult<Option<TrainedHeadInfo>> {
655 + let db = self.db.lock();
656 + Ok(audiofiles_core::analysis::trained_head::train_and_save(&db)?.map(|h| TrainedHeadInfo {
657 + class_count: h.classes.len(),
658 + exemplar_count: h.exemplar_count,
659 + trained_at: h.trained_at,
660 + }))
661 + }
662 +
663 + fn classifier_head_info(&self) -> BackendResult<Option<TrainedHeadInfo>> {
664 + let db = self.db.lock();
665 + Ok(audiofiles_core::analysis::trained_head::load(&db)?.map(|h| TrainedHeadInfo {
666 + class_count: h.classes.len(),
667 + exemplar_count: h.exemplar_count,
668 + trained_at: h.trained_at,
669 + }))
670 + }
671 +
672 + fn clear_classifier_head(&self) -> BackendResult<()> {
673 + let db = self.db.lock();
674 + Ok(audiofiles_core::analysis::trained_head::clear(&db)?)
675 + }
676 +
677 + fn classifier_head_readiness(&self) -> BackendResult<(usize, bool)> {
678 + use audiofiles_core::analysis::trained_head;
642 679 let db = self.db.lock();
643 - Ok(audiofiles_core::analysis::exemplar::auto_apply_library(&db, k)?)
680 + let n = trained_head::labeled_exemplar_count(&db)?;
681 + Ok((n, n >= trained_head::RECOMMENDED_MIN_EXEMPLARS))
644 682 }
645 683
646 684 fn accept_ml_tag(&self, hash: &str, tag: &str) -> BackendResult<()> {
M crates/audiofiles-browser/src/backend/mod.rs +28
@@ -184,6 +184,18 @@ pub struct StorageStats {
184 184 /// only reported, leaving the encoder's clamp as the disclosed last resort.
185 185 pub const FORGE_AUTO_TRIM_OVERSHOOT_KEY: &str = "forge.auto_trim_overshoot";
186 186
187 + /// Lightweight summary of a persisted trained head for the Settings UI (the model's
188 + /// weights stay in core).
189 + #[derive(Debug, Clone, Copy)]
190 + pub struct TrainedHeadInfo {
191 + /// Number of tags the model can score.
192 + pub class_count: usize,
193 + /// Labeled exemplars the model was trained on.
194 + pub exemplar_count: usize,
195 + /// Unix seconds when it was trained.
196 + pub trained_at: i64,
197 + }
198 +
187 199 /// The core abstraction separating UI from data access.
188 200 ///
189 201 /// Every method is synchronous and blocking. The trait is `Send + Sync` so it
@@ -416,6 +428,22 @@ pub trait Backend: Send + Sync {
416 428 &self,
417 429 ) -> BackendResult<Vec<(String, audiofiles_core::analysis::exemplar::TagPolicy)>>;
418 430
431 + // --- Optional trained head (Layer B speed-up for large libraries) ---
432 +
433 + /// Distill the tagged library into a compact per-library model and persist it.
434 + /// Returns a summary, or `None` if there was nothing worth training.
435 + fn train_classifier_head(&self) -> BackendResult<Option<TrainedHeadInfo>>;
436 +
437 + /// Summary of the persisted trained head, if any (and current).
438 + fn classifier_head_info(&self) -> BackendResult<Option<TrainedHeadInfo>>;
439 +
440 + /// Discard the persisted trained head (library-wide auto-tagging falls back to k-NN).
441 + fn clear_classifier_head(&self) -> BackendResult<()>;
442 +
443 + /// `(labeled exemplar count, is a trained head worthwhile at this size)` — drives the
444 + /// advisory hint in the UI.
445 + fn classifier_head_readiness(&self) -> BackendResult<(usize, bool)>;
446 +
419 447 /// Cluster the library's feature vectors into `k` groups (for naming/seeding).
420 448 fn cluster_library(
421 449 &self,
M crates/audiofiles-browser/src/state/classifier.rs +46
@@ -195,6 +195,52 @@ impl BrowserState {
195 195 }
196 196 }
197 197
198 + // ── Optional trained head (Layer B speed-up) ──
199 +
200 + /// Load head summary + readiness once when the section opens.
201 + pub fn ensure_head_loaded(&mut self) {
202 + if !self.classifier.head_loaded {
203 + self.refresh_head();
204 + }
205 + }
206 +
207 + pub fn refresh_head(&mut self) {
208 + self.classifier.head_info = self.backend.classifier_head_info().unwrap_or_default();
209 + self.classifier.head_readiness = self.backend.classifier_head_readiness().ok();
210 + self.classifier.head_loaded = true;
211 + }
212 +
213 + /// Distill the tagged library into a trained head and persist it. Once present, the
214 + /// library-wide "Suggest tags" pass routes through it instead of brute-force k-NN.
215 + pub fn classifier_train_head(&mut self) {
216 + match self.backend.train_classifier_head() {
217 + Ok(Some(info)) => {
218 + self.status = format!(
219 + "Trained model on {} sample{} across {} tag{}",
220 + info.exemplar_count,
221 + if info.exemplar_count == 1 { "" } else { "s" },
222 + info.class_count,
223 + if info.class_count == 1 { "" } else { "s" },
224 + );
225 + }
226 + Ok(None) => {
227 + self.status =
228 + "Not enough tagged samples to train a model yet — tag more first.".to_string();
229 + }
230 + Err(e) => self.status = format!("Training failed: {e}"),
231 + }
232 + self.refresh_head();
233 + }
234 +
235 + /// Discard the trained head; auto-tagging reverts to k-NN.
236 + pub fn classifier_clear_head(&mut self) {
237 + match self.backend.clear_classifier_head() {
238 + Ok(()) => self.status = "Trained model cleared — using nearest-neighbor matching.".to_string(),
239 + Err(e) => self.status = format!("Could not clear model: {e}"),
240 + }
241 + self.refresh_head();
242 + }
243 +
198 244 pub fn ensure_policies_loaded(&mut self) {
199 245 if !self.classifier.policies_loaded {
200 246 self.refresh_policies();
M crates/audiofiles-browser/src/state/tests.rs +63
@@ -1885,4 +1885,67 @@ mod misc {
1885 1885 .iter()
1886 1886 .any(|(t, _)| t == "instrument.drum.kick"));
1887 1887 }
1888 +
1889 + /// Insert a sample with a feature vector and an optional tag (direct DB write,
1890 + /// mirroring `insert_fake_sample`). Used to seed the trained-head tests.
1891 + fn insert_featured_sample(state: &BrowserState, hash: &str, fill: f64, tag: Option<&str>) {
1892 + use audiofiles_core::analysis::classify::{FEATURE_VERSION, NUM_FEATURES};
1893 + insert_fake_sample(state, hash);
1894 + let db = audiofiles_core::db::Database::open(state.data_dir.join("audiofiles.db")).unwrap();
1895 + let vec = vec![fill; NUM_FEATURES];
1896 + let json = serde_json::to_string(&vec).unwrap();
1897 + db.conn()
1898 + .execute(
1899 + "INSERT INTO sample_features (hash, feat_version, vector, computed_at) VALUES (?1, ?2, ?3, 0)",
1900 + rusqlite::params![hash, FEATURE_VERSION, json],
1901 + )
1902 + .unwrap();
1903 + if let Some(t) = tag {
1904 + audiofiles_core::tags::add_tag(&db, hash, t).unwrap();
1905 + }
1906 + }
1907 +
1908 + #[test]
1909 + fn classifier_train_head_and_route() {
1910 + let (mut state, _dir) = make_state();
1911 + // Two separated classes, enough positives each to earn a model.
1912 + for i in 0..6 {
1913 + let k = 0.01 * i as f64;
1914 + insert_featured_sample(&state, &format!("k{i}"), k, Some("instrument.drum.kick"));
1915 + insert_featured_sample(&state, &format!("s{i}"), 10.0 + k, Some("instrument.drum.snare"));
1916 + }
1917 + // An unlabeled sample in the kick cluster.
1918 + insert_featured_sample(&state, "q1", 0.03, None);
1919 +
1920 + // No head yet.
1921 + state.refresh_head();
1922 + assert!(state.classifier.head_info.is_none());
1923 + assert!(state.classifier.head_readiness.is_some());
1924 +
1925 + // Train: head appears, covering both tags.
1926 + state.classifier_train_head();
1927 + let info = state.classifier.head_info.expect("head trained");
1928 + assert_eq!(info.class_count, 2);
1929 + assert_eq!(info.exemplar_count, 12);
1930 +
1931 + // Library-wide auto-tagging now routes through the head and tags the query.
1932 + state.classifier_auto_apply_ml();
1933 + let tags = state.backend.get_sample_tags("q1").unwrap();
1934 + assert!(tags.contains(&"instrument.drum.kick".to_string()), "got {tags:?}");
1935 +
1936 + // Clearing removes the head; routing reverts to k-NN.
1937 + state.classifier_clear_head();
1938 + assert!(state.classifier.head_info.is_none());
1939 + }
1940 +
1941 + #[test]
1942 + fn classifier_train_head_needs_examples() {
1943 + let (mut state, _dir) = make_state();
1944 + // Only two of a tag — under MIN_POSITIVE_EXAMPLES.
1945 + insert_featured_sample(&state, "k0", 0.0, Some("x.kick"));
1946 + insert_featured_sample(&state, "k1", 0.1, Some("x.kick"));
1947 + state.classifier_train_head();
1948 + assert!(state.classifier.head_info.is_none());
1949 + assert!(state.status.contains("Not enough"));
1950 + }
1888 1951 }
M crates/audiofiles-browser/src/state/ui.rs +8
@@ -300,6 +300,14 @@ pub struct ClassifierUiState {
300 300 /// Tag input for adding a new per-tag policy.
301 301 pub new_policy_tag: String,
302 302
303 + // --- Optional trained head (Layer B speed-up) ---
304 + /// Summary of the persisted trained head, if any. `None` = no/stale head.
305 + pub head_info: Option<crate::backend::TrainedHeadInfo>,
306 + /// `(labeled exemplar count, worthwhile)` advisory, loaded with `head_info`.
307 + pub head_readiness: Option<(usize, bool)>,
308 + /// True once head state has been loaded at least once this session.
309 + pub head_loaded: bool,
310 +
303 311 // --- Clustering bootstrap ---
304 312 /// Requested cluster count; 0 = auto (suggest_k).
305 313 pub cluster_k: u32,
M crates/audiofiles-browser/src/ui/classifier.rs +70
@@ -471,6 +471,9 @@ fn draw_autotag_header(ui: &mut egui::Ui, state: &mut BrowserState) {
471 471 }
472 472
473 473 ui.add_space(theme::space::MD);
474 + draw_trained_head(ui, state);
475 +
476 + ui.add_space(theme::space::MD);
474 477 widgets::subsection_label(ui, "Per-tag thresholds");
475 478 ui.label(
476 479 egui::RichText::new("review = surface for review \u{00b7} auto = apply automatically")
@@ -519,6 +522,73 @@ fn draw_autotag_header(ui: &mut egui::Ui, state: &mut BrowserState) {
519 522 });
520 523 }
521 524
525 + /// "Trained model" subsection: the optional distilled head that speeds up the
526 + /// library-wide auto-tagging pass on large libraries. Optional — without it, auto-tagging
527 + /// uses nearest-neighbor matching directly.
528 + fn draw_trained_head(ui: &mut egui::Ui, state: &mut BrowserState) {
529 + state.ensure_head_loaded();
530 + widgets::subsection_label(ui, "Trained model (optional)");
531 + ui.label(
532 + egui::RichText::new(
533 + "For large libraries, distil your tagged samples into a compact model so \
534 + auto-tagging runs much faster. Optional — auto-tagging works without it.",
535 + )
536 + .small()
537 + .color(theme::content_muted()),
538 + );
539 +
540 + match state.classifier.head_info {
541 + Some(info) => {
542 + ui.label(
543 + egui::RichText::new(format!(
544 + "Active — {} tag{} from {} sample{}.",
545 + info.class_count,
546 + if info.class_count == 1 { "" } else { "s" },
547 + info.exemplar_count,
548 + if info.exemplar_count == 1 { "" } else { "s" },
549 + ))
550 + .small(),
551 + );
552 + }
553 + None => {
554 + // Surface the size hint only when there's no model yet.
555 + if let Some((n, worthwhile)) = state.classifier.head_readiness {
556 + let hint = if worthwhile {
557 + format!("No model yet. With {n} tagged samples, training is recommended.")
558 + } else {
559 + format!(
560 + "No model — nearest-neighbor matching is fast enough at {n} tagged \
561 + sample{}.",
562 + if n == 1 { "" } else { "s" }
563 + )
564 + };
565 + ui.label(egui::RichText::new(hint).small().color(theme::content_muted()));
566 + }
567 + }
568 + }
569 +
570 + ui.add_space(theme::space::SM);
571 + ui.horizontal(|ui| {
572 + let train_label =
573 + if state.classifier.head_info.is_some() { "Retrain model" } else { "Train model" };
574 + if ui
575 + .button(train_label)
576 + .on_hover_text("Distil your tagged library into a compact per-library classifier")
577 + .clicked()
578 + {
579 + state.classifier_train_head();
580 + }
581 + if state.classifier.head_info.is_some()
582 + && ui
583 + .button("Clear")
584 + .on_hover_text("Remove the model; auto-tagging reverts to nearest-neighbor matching")
585 + .clicked()
586 + {
587 + state.classifier_clear_head();
588 + }
589 + });
590 + }
591 +
522 592 /// "Clustering" section: unsupervised cold-start grouping.
523 593 fn draw_cluster_header(ui: &mut egui::Ui, state: &mut BrowserState) {
524 594 egui::CollapsingHeader::new(egui::RichText::new("Clustering").strong())
M crates/audiofiles-core/src/analysis/exemplar.rs +20 -4
@@ -218,7 +218,7 @@ fn standardization_params(raw: &[(String, Vec<f64>)]) -> (Vec<f64>, Vec<f64>) {
218 218 (means, stds)
219 219 }
220 220
221 - fn standardize_vec(v: &[f64], means: &[f64], stds: &[f64]) -> Vec<f64> {
221 + pub(crate) fn standardize_vec(v: &[f64], means: &[f64], stds: &[f64]) -> Vec<f64> {
222 222 v.iter()
223 223 .zip(means)
224 224 .zip(stds)
@@ -239,6 +239,22 @@ impl ExemplarIndex {
239 239 self.exemplars.is_empty()
240 240 }
241 241
242 + /// Standardization mean per feature dimension (over the exemplar set).
243 + pub(crate) fn means(&self) -> &[f64] {
244 + &self.means
245 + }
246 +
247 + /// Standardization std-dev per feature dimension (over the exemplar set).
248 + pub(crate) fn stds(&self) -> &[f64] {
249 + &self.stds
250 + }
251 +
252 + /// Iterate `(standardized vector, tags)` over every exemplar — the training set
253 + /// for the distilled head ([`super::trained_head`]).
254 + pub(crate) fn rows(&self) -> impl Iterator<Item = (&[f64], &[String])> {
255 + self.exemplars.iter().map(|e| (e.vector.as_slice(), e.tags.as_slice()))
256 + }
257 +
242 258 /// Score candidate tags for a raw (un-standardized) query vector. `exclude_hash`
243 259 /// drops the query's own exemplar (when scoring an already-labeled sample).
244 260 pub fn score(&self, raw_query: &[f64], k: usize, exclude_hash: Option<&str>) -> Vec<TagScore> {
@@ -304,7 +320,7 @@ impl ExemplarIndex {
304 320 }
305 321
306 322 /// Load a sample's raw feature vector (current version), if present.
307 - fn load_query_vector(db: &Database, hash: &str) -> Result<Option<Vec<f64>>> {
323 + pub(crate) fn load_query_vector(db: &Database, hash: &str) -> Result<Option<Vec<f64>>> {
308 324 let json: Option<String> = db
309 325 .conn()
310 326 .query_row(
@@ -325,7 +341,7 @@ fn load_query_vector(db: &Database, hash: &str) -> Result<Option<Vec<f64>>> {
325 341
326 342 /// Split scored tags into auto-apply / review under each tag's policy, dropping tags the
327 343 /// sample already has and anything below its review threshold.
328 - fn apply_policy(
344 + pub(crate) fn apply_policy(
329 345 db: &Database,
330 346 scores: Vec<TagScore>,
331 347 existing: &std::collections::HashSet<String>,
@@ -355,7 +371,7 @@ fn apply_policy(
355 371 Ok(outcome)
356 372 }
357 373
358 - fn sample_tag_set(db: &Database, hash: &str) -> Result<std::collections::HashSet<String>> {
374 + pub(crate) fn sample_tag_set(db: &Database, hash: &str) -> Result<std::collections::HashSet<String>> {
359 375 Ok(crate::tags::get_sample_tags(db, hash)?.into_iter().collect())
360 376 }
361 377
M crates/audiofiles-core/src/analysis/mod.rs +1
@@ -17,6 +17,7 @@ pub mod loudness;
17 17 pub mod mfcc;
18 18 pub mod spectral;
19 19 pub mod suggest;
20 + pub mod trained_head;
20 21 pub mod waveform;
21 22 pub mod worker;
22 23
A crates/audiofiles-core/src/analysis/trained_head.rs +434
@@ -0,0 +1,434 @@
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::{unix_now, CoreError, Result};
20 + use tracing::instrument;
21 +
22 + use super::classify::{FEATURE_VERSION, NUM_FEATURES};
23 + use super::exemplar::{
24 + self, apply_policy, load_query_vector, sample_tag_set, standardize_vec, MlOutcome, TagScore,
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 { tag: m.tag.clone(), score: sigmoid(z), neighbors: Vec::new() }
84 + })
85 + .collect();
86 + scores.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap().then(a.tag.cmp(&b.tag)));
87 + scores
88 + }
89 + }
90 +
91 + fn sigmoid(z: f64) -> f64 {
92 + 1.0 / (1.0 + (-z).exp())
93 + }
94 +
95 + fn dot(a: &[f64], b: &[f64]) -> f64 {
96 + a.iter().zip(b).map(|(x, y)| x * y).sum()
97 + }
98 +
99 + /// Fit one balanced binary logistic regression by deterministic full-batch gradient
100 + /// descent with L2 regularization. Per-sample class weights counter the heavy
101 + /// positive/negative imbalance of multi-label tags (a tag is rare among all exemplars).
102 + fn train_one(xs: &[&[f64]], labels: &[bool], dims: usize) -> (Vec<f64>, f64) {
103 + let n = xs.len() as f64;
104 + let pos = labels.iter().filter(|y| **y).count() as f64;
105 + let neg = n - pos;
106 + // Balance so positives and negatives contribute equal total weight.
107 + let w_pos = if pos > 0.0 { n / (2.0 * pos) } else { 0.0 };
108 + let w_neg = if neg > 0.0 { n / (2.0 * neg) } else { 0.0 };
109 +
110 + let mut w = vec![0.0; dims];
111 + let mut b = 0.0;
112 + for _ in 0..ITERATIONS {
113 + let mut grad_w = vec![0.0; dims];
114 + let mut grad_b = 0.0;
115 + for (x, &y) in xs.iter().zip(labels) {
116 + let p = sigmoid(dot(&w, x) + b);
117 + let sw = if y { w_pos } else { w_neg };
118 + let err = sw * (p - if y { 1.0 } else { 0.0 });
119 + for d in 0..dims {
120 + grad_w[d] += err * x[d];
121 + }
122 + grad_b += err;
123 + }
124 + for d in 0..dims {
125 + // Mean gradient (weights sum to n) + L2 pull toward zero.
126 + w[d] -= LEARNING_RATE * (grad_w[d] / n + L2 * w[d]);
127 + }
128 + b -= LEARNING_RATE * (grad_b / n);
129 + }
130 + (w, b)
131 + }
132 +
133 + /// Distill the current exemplar store into a trained head **without persisting it**.
134 + /// Returns `None` when there is nothing worth training (no labeled exemplars, or no tag
135 + /// clears [`MIN_POSITIVE_EXAMPLES`]).
136 + #[instrument(skip_all)]
137 + pub fn train(db: &Database) -> Result<Option<TrainedHead>> {
138 + let index = exemplar::build_index(db)?;
139 + if index.is_empty() {
140 + return Ok(None);
141 + }
142 +
143 + // Shared standardized design matrix + per-row tag sets (borrow the index once).
144 + let xs: Vec<&[f64]> = index.rows().map(|(v, _)| v).collect();
145 + let row_tags: Vec<&[String]> = index.rows().map(|(_, t)| t).collect();
146 +
147 + // Tags with enough positives to model; sorted for deterministic class order.
148 + let mut counts: HashMap<&str, usize> = HashMap::new();
149 + for tags in &row_tags {
150 + for t in *tags {
151 + *counts.entry(t.as_str()).or_default() += 1;
152 + }
153 + }
154 + let mut candidates: Vec<&str> = counts
155 + .iter()
156 + .filter(|(_, c)| **c >= MIN_POSITIVE_EXAMPLES)
157 + .map(|(t, _)| *t)
158 + .collect();
159 + candidates.sort_unstable();
160 + if candidates.is_empty() {
161 + return Ok(None);
162 + }
163 +
164 + let classes = candidates
165 + .iter()
166 + .map(|&tag| {
167 + let labels: Vec<bool> =
168 + row_tags.iter().map(|tags| tags.iter().any(|t| t == tag)).collect();
169 + let (weights, bias) = train_one(&xs, &labels, NUM_FEATURES);
170 + TagModel { tag: tag.to_string(), weights, bias }
171 + })
172 + .collect();
173 +
174 + Ok(Some(TrainedHead {
175 + feat_version: FEATURE_VERSION,
176 + means: index.means().to_vec(),
177 + stds: index.stds().to_vec(),
178 + classes,
179 + exemplar_count: index.len(),
180 + trained_at: unix_now(),
181 + }))
182 + }
183 +
184 + /// Distill and persist the head (the background-worker entry point). Persists the fresh
185 + /// model, or [`clear`]s the stored head when there is nothing to train. Returns what was
186 + /// stored (or `None`).
187 + #[instrument(skip_all)]
188 + pub fn train_and_save(db: &Database) -> Result<Option<TrainedHead>> {
189 + match train(db)? {
190 + Some(head) => {
191 + save(db, &head)?;
192 + Ok(Some(head))
193 + }
194 + None => {
195 + clear(db)?;
196 + Ok(None)
197 + }
198 + }
199 + }
200 +
201 + /// Persist the singleton trained head, replacing any prior model.
202 + #[instrument(skip_all)]
203 + pub fn save(db: &Database, head: &TrainedHead) -> Result<()> {
204 + let json = serde_json::to_string(head).map_err(|e| CoreError::Serialization(e.to_string()))?;
205 + db.conn().execute(
206 + "INSERT INTO trained_head (id, feat_version, exemplar_count, model, trained_at)
207 + VALUES (1, ?1, ?2, ?3, ?4)
208 + ON CONFLICT(id) DO UPDATE SET
209 + feat_version = ?1, exemplar_count = ?2, model = ?3, trained_at = ?4",
210 + rusqlite::params![head.feat_version, head.exemplar_count as i64, json, head.trained_at],
211 + )?;
212 + Ok(())
213 + }
214 +
215 + /// Load the stored head, if any. A model trained under a stale `FEATURE_VERSION` is
216 + /// dropped (returns `None`) — callers should retrain.
217 + #[instrument(skip_all)]
218 + pub fn load(db: &Database) -> Result<Option<TrainedHead>> {
219 + let json: Option<String> = db
220 + .conn()
221 + .query_row("SELECT model FROM trained_head WHERE id = 1", [], |row| row.get(0))
222 + .ok();
223 + let Some(json) = json else { return Ok(None) };
224 + let head: TrainedHead =
225 + serde_json::from_str(&json).map_err(|e| CoreError::Serialization(e.to_string()))?;
226 + if !head.is_current() {
227 + return Ok(None);
228 + }
229 + Ok(Some(head))
230 + }
231 +
232 + /// Remove the stored head.
233 + #[instrument(skip_all)]
234 + pub fn clear(db: &Database) -> Result<()> {
235 + db.conn().execute("DELETE FROM trained_head WHERE id = 1", [])?;
236 + Ok(())
237 + }
238 +
239 + /// Count current-version labeled exemplars cheaply (drives the worthwhile/staleness
240 + /// decision without building the full index).
241 + #[instrument(skip_all)]
242 + pub fn labeled_exemplar_count(db: &Database) -> Result<usize> {
243 + let n: i64 = db.conn().query_row(
244 + "SELECT COUNT(DISTINCT t.sample_hash) FROM tags t
245 + JOIN sample_features f ON f.hash = t.sample_hash AND f.feat_version = ?1
246 + WHERE NOT EXISTS (
247 + SELECT 1 FROM tag_provenance p
248 + WHERE p.sample_hash = t.sample_hash AND p.tag = t.tag AND p.source = 'ml')",
249 + [FEATURE_VERSION],
250 + |row| row.get(0),
251 + )?;
252 + Ok(n as usize)
253 + }
254 +
255 + /// Whether a trained head is worth building for the current library size.
256 + #[instrument(skip_all)]
257 + pub fn is_worthwhile(db: &Database) -> Result<bool> {
258 + Ok(labeled_exemplar_count(db)? >= RECOMMENDED_MIN_EXEMPLARS)
259 + }
260 +
261 + /// Score a stored sample against the head under each tag's policy, writing nothing.
262 + /// Mirrors [`exemplar::preview_sample`] but via the distilled model.
263 + #[instrument(skip_all)]
264 + pub fn preview_sample(db: &Database, hash: &str, head: &TrainedHead) -> Result<MlOutcome> {
265 + let Some(raw) = load_query_vector(db, hash)? else {
266 + return Ok(MlOutcome::default());
267 + };
268 + let scores = head.score(&raw);
269 + let existing = sample_tag_set(db, hash)?;
270 + apply_policy(db, scores, &existing)
271 + }
272 +
273 + /// Score a stored sample via the head and **auto-apply** above-threshold tags
274 + /// (`source = 'ml'`), returning what was applied and what's pending review.
275 + #[instrument(skip_all)]
276 + pub fn apply_ml_suggestions(db: &Database, hash: &str, head: &TrainedHead) -> Result<MlOutcome> {
277 + let outcome = preview_sample(db, hash, head)?;
278 + for s in &outcome.auto_applied {
279 + crate::rules::apply_tag_sourced(db, hash, &s.tag, "ml")?;
280 + }
281 + Ok(outcome)
282 + }
283 +
284 + /// Auto-apply above-threshold head tags to every analyzed, non-deleted sample, returning
285 + /// the total tags applied. The head-routed counterpart to [`exemplar::auto_apply_library`]
286 + /// — O(tags) per sample, independent of library size.
287 + #[instrument(skip_all)]
288 + pub fn auto_apply_library(db: &Database, head: &TrainedHead) -> Result<usize> {
289 + let hashes: Vec<String> = {
290 + let mut stmt = db.conn().prepare(
291 + "SELECT s.hash FROM samples s
292 + JOIN sample_features f ON f.hash = s.hash AND f.feat_version = ?1
293 + WHERE s.deleted_at IS NULL",
294 + )?;
295 + stmt.query_map([FEATURE_VERSION], |row| row.get(0))?
296 + .collect::<std::result::Result<Vec<_>, _>>()?
297 + };
298 + let mut applied = 0;
299 + for hash in hashes {
300 + applied += apply_ml_suggestions(db, &hash, head)?.auto_applied.len();
301 + }
302 + Ok(applied)
303 + }
304 +
305 + #[cfg(test)]
306 + mod tests {
307 + use super::*;
308 +
309 + fn insert(db: &Database, hash: &str, vec: [f64; NUM_FEATURES], tags: &[&str]) {
310 + db.conn()
311 + .execute(
312 + "INSERT INTO samples (hash, original_name, file_extension, file_size, import_date, last_modified) \
313 + VALUES (?1, ?1, 'wav', 1, 0, 0)",
314 + [hash],
315 + )
316 + .unwrap();
317 + let json = serde_json::to_string(&vec.to_vec()).unwrap();
318 + db.conn()
319 + .execute(
320 + "INSERT INTO sample_features (hash, feat_version, vector, computed_at) VALUES (?1, ?2, ?3, 0)",
321 + rusqlite::params![hash, FEATURE_VERSION, json],
322 + )
323 + .unwrap();
324 + for t in tags {
325 + crate::tags::add_tag(db, hash, t).unwrap();
326 + }
327 + }
328 +
329 + /// Build two separated clusters of kicks and snares with `n` each, varied per index so
330 + /// the standardizer sees non-zero variance.
331 + fn seed_two_classes(db: &Database, n: usize) {
332 + for i in 0..n {
333 + let k = 0.01 * i as f64;
334 + insert(db, &format!("k{i}"), [k; NUM_FEATURES], &["instrument.drum.kick"]);
335 + insert(db, &format!("s{i}"), [10.0 + k; NUM_FEATURES], &["instrument.drum.snare"]);
336 + }
337 + }
338 +
339 + #[test]
340 + fn no_exemplars_trains_nothing() {
341 + let db = Database::open_in_memory().unwrap();
342 + assert!(train(&db).unwrap().is_none());
343 + }
344 +
345 + #[test]
346 + fn too_few_positives_trains_nothing() {
347 + let db = Database::open_in_memory().unwrap();
348 + // Only 2 of each tag (< MIN_POSITIVE_EXAMPLES).
349 + seed_two_classes(&db, 2);
350 + assert!(train(&db).unwrap().is_none());
351 + }
352 +
353 + #[test]
354 + fn distinguishes_two_classes() {
355 + let db = Database::open_in_memory().unwrap();
356 + seed_two_classes(&db, 8);
357 + let head = train(&db).unwrap().expect("should train");
358 + assert_eq!(head.classes.len(), 2);
359 + assert!(head.is_current());
360 +
361 + // A query in the kick cluster scores kick above snare.
362 + let near_kick = [0.04; NUM_FEATURES];
363 + let scores = head.score(&near_kick);
364 + let kick = scores.iter().find(|s| s.tag == "instrument.drum.kick").unwrap();
365 + let snare = scores.iter().find(|s| s.tag == "instrument.drum.snare").unwrap();
366 + assert!(kick.score > snare.score, "kick {} vs snare {}", kick.score, snare.score);
367 + assert!(kick.score > 0.5);
368 + }
369 +
370 + #[test]
371 + fn save_load_round_trip_and_clear() {
372 + let db = Database::open_in_memory().unwrap();
373 + seed_two_classes(&db, 6);
374 + let head = train_and_save(&db).unwrap().unwrap();
375 + let loaded = load(&db).unwrap().expect("persisted");
376 + assert_eq!(loaded.classes.len(), head.classes.len());
377 + assert_eq!(loaded.exemplar_count, head.exemplar_count);
378 + clear(&db).unwrap();
379 + assert!(load(&db).unwrap().is_none());
380 + }
381 +
382 + #[test]
383 + fn stale_feature_version_is_ignored_on_load() {
384 + let db = Database::open_in_memory().unwrap();
385 + seed_two_classes(&db, 6);
386 + let mut head = train(&db).unwrap().unwrap();
387 + head.feat_version = FEATURE_VERSION + 1;
388 + save(&db, &head).unwrap();
389 + // Stored row exists but is stale → load returns None.
390 + assert!(load(&db).unwrap().is_none());
391 + }
392 +
393 + #[test]
394 + fn deterministic_across_runs() {
395 + let train_once = || {
396 + let db = Database::open_in_memory().unwrap();
397 + seed_two_classes(&db, 7);
398 + train(&db).unwrap().unwrap().classes
399 + };
400 + let a = train_once();
401 + let b = train_once();
402 + assert_eq!(a.len(), b.len());
403 + for (ca, cb) in a.iter().zip(&b) {
404 + assert_eq!(ca.tag, cb.tag);
405 + assert_eq!(ca.bias.to_bits(), cb.bias.to_bits(), "bias differs for {}", ca.tag);
406 + for (wa, wb) in ca.weights.iter().zip(&cb.weights) {
407 + assert_eq!(wa.to_bits(), wb.to_bits());
408 + }
409 + }
410 + }
411 +
412 + #[test]
413 + fn apply_auto_tags_via_head() {
414 + let db = Database::open_in_memory().unwrap();
415 + seed_two_classes(&db, 10);
416 + let head = train(&db).unwrap().unwrap();
417 + // Unlabeled query in the kick cluster.
418 + insert(&db, "q", [0.05; NUM_FEATURES], &[]);
419 + let outcome = apply_ml_suggestions(&db, "q", &head).unwrap();
420 + let applied: Vec<&str> = outcome.auto_applied.iter().map(|s| s.tag.as_str()).collect();
421 + assert!(applied.contains(&"instrument.drum.kick"), "got {applied:?}");
422 + // Provenance is 'ml'.
423 + let prov = crate::rules::sample_tag_provenance(&db, "q").unwrap();
424 + assert!(prov.iter().any(|(t, s, _)| t == "instrument.drum.kick" && s == "ml"));
425 + }
426 +
427 + #[test]
428 + fn labeled_count_and_worthwhile() {
429 + let db = Database::open_in_memory().unwrap();
430 + seed_two_classes(&db, 5); // 10 labeled exemplars
431 + assert_eq!(labeled_exemplar_count(&db).unwrap(), 10);
432 + assert!(!is_worthwhile(&db).unwrap()); // well under RECOMMENDED_MIN_EXEMPLARS
433 + }
434 + }
M crates/audiofiles-core/src/db.rs +22 -6
@@ -1296,6 +1296,20 @@ BEGIN
1296 1296 END;
1297 1297 "#;
1298 1298
1299 + const MIGRATION_024: &str = r#"
1300 + -- Phase 6: optional trained head (per-library distilled logistic-regression classifier).
1301 + -- A local, regenerable cache derived entirely from sample_features + tags; NOT synced
1302 + -- (no changelog triggers) because it rebuilds from the exemplar store on any device.
1303 + -- Singleton row (id = 1); the whole model travels as one JSON blob.
1304 + CREATE TABLE IF NOT EXISTS trained_head (
1305 + id INTEGER PRIMARY KEY CHECK (id = 1),
1306 + feat_version INTEGER NOT NULL,
1307 + exemplar_count INTEGER NOT NULL,
1308 + model TEXT NOT NULL,
1309 + trained_at INTEGER NOT NULL
1310 + );
1311 + "#;
1312 +
1299 1313 /// Register `hash_row_id(salt, key) -> TEXT` as a deterministic SQLite
1300 1314 /// function on the given connection. Used by the M018 sync triggers so the
1301 1315 /// `sync_changelog.row_id` field never carries cleartext content (tag strings,
@@ -1399,6 +1413,7 @@ impl Database {
1399 1413 MIGRATION_021,
1400 1414 MIGRATION_022,
1401 1415 MIGRATION_023,
1416 + MIGRATION_024,
1402 1417 ];
1403 1418
1404 1419 for (i, sql) in MIGRATIONS.iter().enumerate() {
@@ -1564,6 +1579,7 @@ mod tests {
1564 1579 "tag_provenance",
1565 1580 "tag_rules",
1566 1581 "tags",
1582 + "trained_head",
1567 1583 "user_config",
1568 1584 "vfs",
1569 1585 "vfs_nodes",
@@ -1579,7 +1595,7 @@ mod tests {
1579 1595 .conn()
1580 1596 .query_row("PRAGMA user_version", [], |row| row.get(0))
1581 1597 .unwrap();
1582 - assert_eq!(version, 23);
1598 + assert_eq!(version, 24);
1583 1599 }
1584 1600
1585 1601 #[test]
@@ -1590,7 +1606,7 @@ mod tests {
1590 1606 .conn()
1591 1607 .query_row("PRAGMA user_version", [], |row| row.get(0))
1592 1608 .unwrap();
1593 - assert_eq!(version, 23);
1609 + assert_eq!(version, 24);
1594 1610 }
1595 1611
1596 1612 /// Open a fresh file-backed DB, close, reopen. The second open re-enters
@@ -1609,7 +1625,7 @@ mod tests {
1609 1625 .conn()
1610 1626 .query_row("PRAGMA user_version", [], |row| row.get(0))
1611 1627 .unwrap();
1612 - assert_eq!(version, 23);
1628 + assert_eq!(version, 24);
1613 1629 }
1614 1630
1615 1631 /// Simulates the worst-case recovery path: a prior partial migration left
@@ -1653,7 +1669,7 @@ mod tests {
1653 1669 .conn()
1654 1670 .query_row("PRAGMA user_version", [], |row| row.get(0))
1655 1671 .unwrap();
1656 - assert_eq!(version, 23);
1672 + assert_eq!(version, 24);
1657 1673 }
1658 1674
1659 1675 /// M018 contract: the `sync_changelog.row_id` for sensitive tables must
@@ -1875,7 +1891,7 @@ mod tests {
1875 1891 let initial_version: i32 = conn
1876 1892 .query_row("PRAGMA user_version", [], |row| row.get(0))
1877 1893 .unwrap();
1878 - assert_eq!(initial_version, 23);
1894 + assert_eq!(initial_version, 24);
1879 1895
1880 1896 let batch = format!(
1881 1897 "BEGIN;\n{}\nPRAGMA user_version = 999;\nCOMMIT;",
@@ -1938,7 +1954,7 @@ mod tests {
1938 1954 .conn()
1939 1955 .query_row("PRAGMA user_version", [], |row| row.get(0))
1940 1956 .unwrap();
1941 - assert_eq!(version, 23);
1957 + assert_eq!(version, 24);
1942 1958 }
1943 1959
1944 1960 #[test]