Skip to main content

max / audiofiles

20.8 KB · 516 lines History Blame Raw
1 //! Scoreable corpus rows, and the in-memory indexes built from them.
2 //!
3 //! Extracted from [`crate::layer_eval`] when [`crate::stability`] arrived and
4 //! needed the same four things: read the analysed vault back as (vector, labels,
5 //! truth) rows projected onto a label space, split them without an RNG, and build
6 //! a real `ExemplarIndex` from an arbitrary subset. Same reason [`crate::labelled`]
7 //! exists: the step is subtle enough that two copies would drift, and the
8 //! subtleties are all about what is silently excluded.
9 //!
10 //! What lives here is label-space-agnostic and policy-free. The gate belongs to
11 //! `layer_eval`, the bar to `stability`; neither is a property of a row.
12
13 use std::collections::{BTreeSet, HashMap};
14
15 use audiofiles_core::analysis::afcl::{
16 self, Afcl, AfclExemplar, AfclManifest, DEFAULT_IMPORT_WEIGHT,
17 };
18 use audiofiles_core::analysis::features::{FEATURE_VERSION, NUM_FEATURES};
19 use audiofiles_core::db::Database;
20
21 use crate::families::{self, LabelSpace};
22 use crate::labelled;
23
24 /// One corpus sample: its vector, its labels, and the class it is scored against.
25 pub(crate) struct Row {
26 pub(crate) hash: String,
27 pub(crate) vector: Vec<f64>,
28 /// Labels in the evaluation's space, so the index under test carries the tags
29 /// being graded.
30 pub(crate) tags: Vec<String>,
31 /// The single class this row is ground truth for, or `None` when the corpus
32 /// gave it more than one. Content-addressed import collapses a file that
33 /// appears in two class folders into one row carrying both tags; its true
34 /// class is undecidable, so it trains but is never tested.
35 pub(crate) truth: Option<String>,
36 /// The corpus folder(s) this row came from, before projection. Kept so a
37 /// family's members can be broken out by where they came from: `low` is kick
38 /// plus tom, and whether the layer recovers a tom as readily as a kick is the
39 /// open split `af-coarse-families` asks about.
40 pub(crate) origin: String,
41 /// The sample's filename, as imported. Only the filename rules read it, and
42 /// only to answer whether this sample is in the population the layer exists
43 /// to serve: `starter_rules` already labels 97.7% of this corpus off its
44 /// name, and every accuracy number so far was measured on exactly that
45 /// population. See [`crate::stability`].
46 pub(crate) name: String,
47 }
48
49 /// What [`load_rows`] set aside, so no exclusion is silent.
50 pub(crate) struct Dropped {
51 /// Rows whose every corpus tag projects nowhere in this label space.
52 pub(crate) unprojectable: usize,
53 /// Corpus labels those rows came from.
54 pub(crate) origins: BTreeSet<String>,
55 }
56
57 /// Read the analysed corpus back out of the scratch vault as scoreable rows,
58 /// projected onto `space`.
59 ///
60 /// The projection happens here rather than in [`labelled`] on purpose. The vault
61 /// is ground truth at the finest resolution the corpus carries, which is what
62 /// `afcl_gen` exports from; grading at a coarser resolution is a property of the
63 /// evaluation, not of the corpus, so it belongs on the read side. That also keeps
64 /// the exported layer byte-identical and leaves the retired instrument question
65 /// runnable instead of deleted.
66 pub(crate) fn load_rows(db: &Database, space: LabelSpace) -> Result<(Vec<Row>, Dropped), String> {
67 let conn = db.conn();
68
69 let mut tags_by_hash: HashMap<String, Vec<String>> = HashMap::new();
70 {
71 let mut stmt = conn
72 .prepare("SELECT sample_hash, tag FROM tags")
73 .map_err(|e| format!("tags query: {e}"))?;
74 let rows = stmt
75 .query_map([], |row| {
76 Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
77 })
78 .map_err(|e| format!("tags query: {e}"))?;
79 for r in rows {
80 let (hash, tag) = r.map_err(|e| format!("tags row: {e}"))?;
81 tags_by_hash.entry(hash).or_default().push(tag);
82 }
83 }
84
85 // Names are a left join in spirit: a row with no `samples` entry is one this
86 // module fabricated in a test, and it still has to be scoreable.
87 let mut name_by_hash: HashMap<String, String> = HashMap::new();
88 {
89 let mut stmt = conn
90 .prepare("SELECT hash, original_name FROM samples")
91 .map_err(|e| format!("names query: {e}"))?;
92 let rows = stmt
93 .query_map([], |row| {
94 Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
95 })
96 .map_err(|e| format!("names query: {e}"))?;
97 for r in rows {
98 let (hash, name) = r.map_err(|e| format!("names row: {e}"))?;
99 name_by_hash.insert(hash, name);
100 }
101 }
102
103 let mut out = Vec::new();
104 let mut dropped = Dropped {
105 unprojectable: 0,
106 origins: BTreeSet::new(),
107 };
108 let mut stmt = conn
109 .prepare("SELECT hash, vector FROM sample_features WHERE feat_version = ?1 ORDER BY hash")
110 .map_err(|e| format!("features query: {e}"))?;
111 let rows = stmt
112 .query_map([FEATURE_VERSION], |row| {
113 Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
114 })
115 .map_err(|e| format!("features query: {e}"))?;
116 for r in rows {
117 let (hash, json) = r.map_err(|e| format!("features row: {e}"))?;
118 let Some(mut corpus_tags) = tags_by_hash.remove(&hash) else {
119 continue;
120 };
121 corpus_tags.sort();
122 corpus_tags.dedup();
123 let vector: Vec<f64> =
124 serde_json::from_str(&json).map_err(|e| format!("vector for {hash}: {e}"))?;
125 // The index drops these too (`is_usable_vector`), so counting them as
126 // testable would score the layer on samples it never sees.
127 if vector.len() != NUM_FEATURES || !vector.iter().all(|x| x.is_finite()) {
128 continue;
129 }
130
131 // Project, then dedup again: two instrument tags landing in one family is
132 // not ambiguity, it is the coarser question being easier. A file in both
133 // the kick and tom folders has no instrument truth and a perfectly good
134 // family one.
135 let mut tags: Vec<String> = corpus_tags
136 .iter()
137 .filter_map(|t| families::project(space, t))
138 .map(str::to_string)
139 .collect();
140 tags.sort();
141 tags.dedup();
142 if tags.is_empty() {
143 dropped.unprojectable += 1;
144 dropped.origins.extend(
145 corpus_tags
146 .iter()
147 .map(|t| labelled::label_for_tag(t).to_string()),
148 );
149 continue;
150 }
151
152 let truth = (tags.len() == 1).then(|| tags[0].clone());
153 let origin = corpus_tags
154 .iter()
155 .map(|t| labelled::label_for_tag(t))
156 .collect::<Vec<_>>()
157 .join("+");
158 let name = name_by_hash.get(&hash).cloned().unwrap_or_default();
159 out.push(Row {
160 hash,
161 vector,
162 tags,
163 truth,
164 origin,
165 name,
166 });
167 }
168 Ok((out, dropped))
169 }
170
171 /// Assign each row a fold, stratified by class.
172 ///
173 /// Round-robin down each class's hash-sorted list rather than a shuffle: the
174 /// content hash is already an arbitrary order with respect to the source pack a
175 /// file came from, so this needs no RNG and two runs over the same corpus produce
176 /// the same folds. Ambiguous rows get no fold; they train everywhere.
177 pub(crate) fn assign_folds(rows: &[Row], folds: usize) -> Vec<Option<usize>> {
178 let mut seen_per_class: HashMap<&str, usize> = HashMap::new();
179 rows.iter()
180 .map(|r| {
181 let truth = r.truth.as_deref()?;
182 let n = seen_per_class.entry(truth).or_default();
183 let fold = *n % folds;
184 *n += 1;
185 Some(fold)
186 })
187 .collect()
188 }
189
190 /// Build an in-memory vault holding just `local` rows, so the index under test is
191 /// built by the same `build_index` the app calls.
192 ///
193 /// Every exemplar here is `local` (weight 1.0) where the shipped ones arrive from
194 /// an imported layer at a lower weight. That does not change a score: the weight
195 /// is a constant multiplier inside a sum that is then divided by its own total,
196 /// so a uniform weight cancels. It matters only against a mixed index of user
197 /// labels plus the layer, which is [`mixed_db`].
198 pub(crate) fn local_db(local: &[&Row]) -> Result<Database, String> {
199 mixed_db(local, &[])
200 }
201
202 /// Build an in-memory vault holding `local` rows at weight 1.0 and `imported`
203 /// rows as an enabled `.afcl` layer at [`DEFAULT_IMPORT_WEIGHT`].
204 ///
205 /// This is the deployment shape, and it is not a uniform index: the weight only
206 /// cancels when every exemplar carries the same one. Here the ratio is real, and
207 /// so is the second-order effect that makes this worth simulating rather than
208 /// reasoning about — `build_index` fits the standardization means and standard
209 /// deviations over local **and** imported exemplars together, so adding a user's
210 /// labels moves the space every distance is measured in.
211 ///
212 /// The layer goes in through `afcl::import` rather than a hand-written INSERT so
213 /// the weight, the enabled flag and the `feat_version` gate are the ones the app
214 /// applies. Imported exemplars carry vector + tags and no hash, exactly as a
215 /// shipped layer does.
216 pub(crate) fn mixed_db(local: &[&Row], imported: &[&Row]) -> Result<Database, String> {
217 let db = Database::open_in_memory().map_err(|e| format!("open_in_memory: {e}"))?;
218 db.transaction(|_tx| {
219 for row in local {
220 db.conn().execute(
221 "INSERT INTO samples (hash, original_name, file_extension, file_size, import_date, last_modified)
222 VALUES (?1, ?2, 'wav', 1, 0, 0)",
223 rusqlite::params![&row.hash, &row.name],
224 )?;
225 let json = serde_json::to_string(&row.vector).unwrap_or_default();
226 db.conn().execute(
227 "INSERT INTO sample_features (hash, feat_version, vector, computed_at) VALUES (?1, ?2, ?3, 0)",
228 rusqlite::params![&row.hash, FEATURE_VERSION, json],
229 )?;
230 }
231 Ok(())
232 })
233 .map_err(|e| format!("seeding the index: {e}"))?;
234 // Tags outside the transaction: `add_tag` is the public path and manages its
235 // own writes, and this is an in-memory DB where the fsync cost it avoids does
236 // not exist.
237 for row in local {
238 for tag in &row.tags {
239 audiofiles_core::tags::add_tag(&db, &row.hash, tag)
240 .map_err(|e| format!("tagging the index: {e}"))?;
241 }
242 }
243
244 if !imported.is_empty() {
245 let exemplars: Vec<AfclExemplar> = imported
246 .iter()
247 .map(|r| AfclExemplar {
248 vector: r.vector.clone(),
249 tags: r.tags.clone(),
250 })
251 .collect();
252 let layer = Afcl {
253 manifest: AfclManifest {
254 afcl_version: afcl::AFCL_VERSION,
255 feat_version: FEATURE_VERSION,
256 name: "simulated imported layer".into(),
257 description: "built by audiofiles-bench, never written to disk".into(),
258 kind: afcl::LayerKind::Official.as_str().to_string(),
259 created_at: 0,
260 license_note: String::new(),
261 exemplar_count: exemplars.len(),
262 rule_count: 0,
263 policy_count: 0,
264 },
265 exemplars,
266 rules: Vec::new(),
267 policy: Vec::new(),
268 };
269 let summary = afcl::import(&db, &layer, Some(afcl::BUNDLED_SOURCE))
270 .map_err(|e| format!("importing the simulated layer: {e}"))?;
271 // An exemplar dropped at import is one the measurement thinks it has and
272 // does not. `load_rows` already applied the same finiteness filter, so a
273 // mismatch here means the two disagree and the run is not what it says.
274 if summary.exemplars != imported.len() {
275 return Err(format!(
276 "the simulated layer imported {} of {} exemplars",
277 summary.exemplars,
278 imported.len()
279 ));
280 }
281 }
282 Ok(db)
283 }
284
285 /// The weight an imported layer carries beside the user's own labels, restated
286 /// here so a report can print the ratio it was measured at.
287 pub(crate) const IMPORT_WEIGHT: f64 = DEFAULT_IMPORT_WEIGHT;
288
289 #[cfg(test)]
290 mod tests {
291 use super::*;
292
293 fn row(hash: &str, tag: &str) -> Row {
294 Row {
295 hash: hash.into(),
296 vector: vec![0.5; NUM_FEATURES],
297 tags: vec![tag.into()],
298 truth: Some(tag.into()),
299 origin: labelled::label_for_tag(tag).into(),
300 name: format!("{hash}.wav"),
301 }
302 }
303
304 #[test]
305 fn folds_are_stratified_and_ambiguous_rows_get_none() {
306 let mut rows: Vec<Row> = (0..10)
307 .map(|i| row(&format!("k{i}"), "instrument.drum.kick"))
308 .collect();
309 rows.extend((0..4).map(|i| row(&format!("s{i}"), "instrument.drum.snare")));
310 rows.push(Row {
311 hash: "dup".into(),
312 vector: vec![0.0; NUM_FEATURES],
313 tags: Vec::new(),
314 truth: None,
315 origin: String::new(),
316 name: "dup.wav".into(),
317 });
318
319 let folds = assign_folds(&rows, 5);
320 assert_eq!(
321 folds.last().copied().flatten(),
322 None,
323 "ambiguous row trains everywhere"
324 );
325
326 // Every fold holds two of the ten kicks: stratification, not a global
327 // round-robin that would leave a fold without any snares.
328 for f in 0..5 {
329 let kicks = rows
330 .iter()
331 .zip(&folds)
332 .filter(|(r, fold)| {
333 **fold == Some(f) && r.truth.as_deref() == Some("instrument.drum.kick")
334 })
335 .count();
336 assert_eq!(kicks, 2, "fold {f}");
337 }
338 // Four snares over five folds: one fold is short, and none holds two.
339 for f in 0..5 {
340 let snares = rows
341 .iter()
342 .zip(&folds)
343 .filter(|(r, fold)| {
344 **fold == Some(f) && r.truth.as_deref() == Some("instrument.drum.snare")
345 })
346 .count();
347 assert!(snares <= 1, "fold {f} holds {snares} snares");
348 }
349 }
350
351 #[test]
352 fn folds_are_deterministic() {
353 let rows: Vec<Row> = (0..20)
354 .map(|i| row(&format!("k{i}"), "instrument.drum.kick"))
355 .collect();
356 assert_eq!(assign_folds(&rows, 5), assign_folds(&rows, 5));
357 }
358
359 #[test]
360 fn a_fold_index_scores_a_held_out_sample_from_its_neighbours() {
361 // End to end over the real index: two tight clusters, one held out
362 // sample, and the index must answer with the cluster it sits in.
363 let mut rows = Vec::new();
364 for i in 0..5 {
365 let mut r = row(&format!("k{i}"), "instrument.drum.kick");
366 r.vector = vec![0.01 * f64::from(i); NUM_FEATURES];
367 rows.push(r);
368 }
369 for i in 0..5 {
370 let mut r = row(&format!("s{i}"), "instrument.drum.snare");
371 r.vector = vec![100.0 + f64::from(i); NUM_FEATURES];
372 rows.push(r);
373 }
374 let refs: Vec<&Row> = rows.iter().collect();
375 let db = local_db(&refs).unwrap();
376 let index = audiofiles_core::analysis::exemplar::build_index(&db).unwrap();
377 assert_eq!(index.len(), 10);
378
379 let query = vec![0.02; NUM_FEATURES];
380 let scored = index.score(&query, audiofiles_core::analysis::exemplar::DEFAULT_K, None);
381 assert_eq!(
382 scored.first().map(|s| s.tag.as_str()),
383 Some("instrument.drum.kick")
384 );
385 }
386
387 #[test]
388 fn load_rows_marks_a_multi_tagged_sample_ambiguous() {
389 let rows = [
390 row("a", "instrument.drum.kick"),
391 Row {
392 hash: "b".into(),
393 vector: vec![0.5; NUM_FEATURES],
394 tags: vec![
395 "instrument.drum.kick".into(),
396 "instrument.drum.snare".into(),
397 ],
398 truth: None,
399 origin: "kick+snare".into(),
400 name: "b.wav".into(),
401 },
402 ];
403 let refs: Vec<&Row> = rows.iter().collect();
404 let db = local_db(&refs).unwrap();
405
406 let (read, dropped) = load_rows(&db, LabelSpace::Instrument).unwrap();
407 assert_eq!(read.len(), 2);
408 assert_eq!(dropped.unprojectable, 0);
409 let b = read.iter().find(|r| r.hash == "b").unwrap();
410 assert_eq!(b.truth, None, "two class tags means no ground truth");
411 assert_eq!(b.tags.len(), 2, "but it still trains with both labels");
412 }
413
414 #[test]
415 fn the_family_projection_resolves_an_instrument_level_ambiguity() {
416 // A file in both the kick and tom folders has no instrument truth and a
417 // perfectly good family one: the coarser question is the easier question,
418 // and that has to show up as a testable row rather than a dropped one.
419 let rows = [Row {
420 hash: "a".into(),
421 vector: vec![0.5; NUM_FEATURES],
422 tags: vec!["instrument.drum.kick".into(), "instrument.drum.tom".into()],
423 truth: None,
424 origin: "kick+tom".into(),
425 name: "a.wav".into(),
426 }];
427 let refs: Vec<&Row> = rows.iter().collect();
428 let db = local_db(&refs).unwrap();
429
430 let (read, _) = load_rows(&db, LabelSpace::Family).unwrap();
431 assert_eq!(read[0].truth.as_deref(), Some("family.low"));
432 assert_eq!(read[0].origin, "kick+tom");
433 }
434
435 #[test]
436 fn percussion_leaves_the_family_run_entirely() {
437 // Not merely untested: an exemplar carrying a guessed family label would
438 // pull the index toward that guess for every other sample too.
439 let rows = [Row {
440 hash: "p".into(),
441 vector: vec![0.5; NUM_FEATURES],
442 tags: vec!["instrument.percussion".into()],
443 truth: Some("instrument.percussion".into()),
444 origin: "percussion".into(),
445 name: "p.wav".into(),
446 }];
447 let refs: Vec<&Row> = rows.iter().collect();
448 let db = local_db(&refs).unwrap();
449
450 let (read, dropped) = load_rows(&db, LabelSpace::Family).unwrap();
451 assert!(read.is_empty(), "it must not train either");
452 assert_eq!(dropped.unprojectable, 1);
453 assert!(dropped.origins.contains("percussion"));
454 }
455
456 #[test]
457 fn load_rows_carries_the_filename_back() {
458 // The rule-miss population is keyed off this, and a silently empty name
459 // would read as "no filename rule fires" for the whole corpus, which is
460 // the opposite of the truth.
461 let rows = [row("a", "instrument.drum.kick")];
462 let refs: Vec<&Row> = rows.iter().collect();
463 let db = local_db(&refs).unwrap();
464 let (read, _) = load_rows(&db, LabelSpace::Instrument).unwrap();
465 assert_eq!(read[0].name, "a.wav");
466 }
467
468 #[test]
469 fn a_mixed_index_holds_both_populations() {
470 let local: Vec<Row> = (0..3)
471 .map(|i| row(&format!("l{i}"), "instrument.drum.kick"))
472 .collect();
473 let imported: Vec<Row> = (0..4)
474 .map(|i| row(&format!("i{i}"), "instrument.drum.snare"))
475 .collect();
476 let l: Vec<&Row> = local.iter().collect();
477 let i: Vec<&Row> = imported.iter().collect();
478 let db = mixed_db(&l, &i).unwrap();
479
480 let index = audiofiles_core::analysis::exemplar::build_index(&db).unwrap();
481 assert_eq!(index.len(), 7, "local and imported both reach the index");
482
483 let layers = afcl::list_layers(&db).unwrap();
484 assert_eq!(layers.len(), 1);
485 assert!(layers[0].enabled, "a disabled layer contributes nothing");
486 assert!((layers[0].weight - IMPORT_WEIGHT).abs() < 1e-9);
487 }
488
489 #[test]
490 fn an_imported_exemplar_outvoted_by_local_ones_shows_the_weight_is_live() {
491 // The property `mixed_db` exists for. Two exemplars equidistant from the
492 // query, one local and one imported: the local tag must score 1/(1+0.5)
493 // of the neighbourhood rather than half of it. If the weight were being
494 // cancelled away this would come back 0.5.
495 let mut near_local = row("l0", "instrument.drum.kick");
496 near_local.vector = vec![0.0; NUM_FEATURES];
497 let mut near_imported = row("i0", "instrument.drum.snare");
498 near_imported.vector = vec![0.0; NUM_FEATURES];
499
500 let db = mixed_db(&[&near_local], &[&near_imported]).unwrap();
501 let index = audiofiles_core::analysis::exemplar::build_index(&db).unwrap();
502 let scored = index.score(&vec![0.0; NUM_FEATURES], 15, None);
503
504 let kick = scored
505 .iter()
506 .find(|s| s.tag == "instrument.drum.kick")
507 .expect("the local exemplar is in the neighbourhood");
508 assert!(
509 (kick.score - (1.0 / (1.0 + IMPORT_WEIGHT))).abs() < 1e-9,
510 "local should carry {:.3} of the weight, carried {:.3}",
511 1.0 / (1.0 + IMPORT_WEIGHT),
512 kick.score
513 );
514 }
515 }
516