From e8575780e46789863f55061579994007231ad060 Mon Sep 17 00:00:00 2001 From: clawbot <35+clawbot@noreply.example.org> Date: Tue, 22 Sep 2026 20:45:59 +0200 Subject: [PATCH] Content-similarity search surface over the CLIP index (closes #50) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds lib.mldata search over the CLIP index (#49): forFile returns a file's stored payload; similar ranks nearest files by cosine on the CLIP embedding; searchByEmbedding ranks the index against a caller-supplied query vector. All RAM-only, reusing the packed Float32Array index and id list. No text encoder is bundled — the caller provides the query embedding. (Redo of the reverted first attempt, now tsc-clean.) Model: opus-4-8 --- src/library/index.ts | 16 +++-- src/library/mlsearch.ts | 129 ++++++++++++++++++++++++++++++++++ test/library/mlsearch.test.ts | 108 ++++++++++++++++++++++++++++ 3 files changed, 249 insertions(+), 4 deletions(-) create mode 100644 src/library/mlsearch.ts create mode 100644 test/library/mlsearch.test.ts diff --git a/src/library/index.ts b/src/library/index.ts index e8c1089..d458df1 100644 --- a/src/library/index.ts +++ b/src/library/index.ts @@ -48,6 +48,7 @@ import { type EnsureOptions, type EnsureResult, } from "./content.js"; +import { makeMLDataAPI, type MLDataAPI } from "./mlsearch.js"; export { Album, @@ -71,6 +72,7 @@ export { type EnsureResult, type EnsureEvent, } from "./content.js"; +export { type MLDataAPI, type SimilarResult } from "./mlsearch.js"; import type { CollectionsPage, FilesPage } from "../client.js"; import { MLDATA_BATCH_SIZE, type MLData } from "../mldata-fetch.js"; import type { Collection, EnteFile } from "../model/types.js"; @@ -189,6 +191,10 @@ export class Library { // The thumbnail-prefetch surface (issue #46): drives the thumbnail pool // with priority, dedup, and abort. readonly thumbnails: ThumbnailsAPI; + // The content-similarity search surface over the CLIP index (issue #50). + // Present whether or not ML fetching is enabled; with no ML store it + // returns empty results. + readonly mldata: MLDataAPI; private readonly client: LibraryClient; private readonly store: MetadataStore; @@ -200,7 +206,7 @@ export class Library { private readonly onProgress?: RefreshProgressCallback; private readonly pools: RequestPools; // The ML-data cache, present only when the client can fetch ML data. - private readonly mldata?: MLDataStore; + private readonly mlStore?: MLDataStore; private timer?: ReturnType; private refreshing = false; @@ -243,7 +249,7 @@ export class Library { this.intervalMs = args.intervalMs; this.onProgress = args.onProgress; this.pools = args.pools; - this.mldata = args.mldata; + this.mlStore = args.mldata; this.cache = args.cache; this.lastRecords = this.deriveNow(); @@ -265,6 +271,8 @@ export class Library { return this.cache.ensureThumbnails(opts); }, }; + // Reads the ML store live so results grow as ML data is fetched. + this.mldata = makeMLDataAPI(() => this.mlStore); } // Load the cache and start the refresh loop. With an empty cache the first @@ -386,7 +394,7 @@ export class Library { for (const c of collections) { files += this.store.listFiles(c.id).length; } - const ml = this.mldata?.stats(); + const ml = this.mlStore?.stats(); const originals = this.cache?.originalsStatus(); return { userID: this.store.userID, @@ -583,7 +591,7 @@ export class Library { // advanced), through the metadata pool, and update the CLIP index. Guarded // so passes never overlap; a failure is reported, not thrown. private async runMLFetch(): Promise { - const mldata = this.mldata; + const mldata = this.mlStore; // Bind so the call keeps the client as its receiver when invoked // through the pool below. const fetchMLData = this.client.fetchMLData?.bind(this.client); diff --git a/src/library/mlsearch.ts b/src/library/mlsearch.ts new file mode 100644 index 0000000..d23ef26 --- /dev/null +++ b/src/library/mlsearch.ts @@ -0,0 +1,129 @@ +// The content-similarity search surface over the CLIP index (issue #50). +// +// This is `lib.mldata`. It answers three questions against the ML-data cache +// (#49) without touching the network: +// +// - `forFile` returns the whole stored payload (face boxes, landmarks, +// embedding) for a file, read from disk on demand — the only method here +// that touches the disk, and the only one that is async. +// - `similar` and `searchByEmbedding` rank fileIDs by cosine similarity over +// the packed `Float32Array` index alone. That index (~50k×512) already +// lives in RAM, so each query is a plain loop over it and nothing else. +// +// quak bundles no text encoder (owner-deferred), so `searchByEmbedding` takes +// the query vector the caller has produced elsewhere; `similar` uses the +// query file's own indexed embedding. + +import type { MLData } from "../mldata-fetch.js"; +import type { MLDataStore, MLIndex } from "./mldata.js"; + +// How many nearest files a query returns when the caller names no limit. +const DEFAULT_LIMIT = 20; + +// One ranked result: a fileID and its cosine similarity to the query, in +// [-1, 1]. Callers wanting only the ids read `.fileID`. +export interface SimilarResult { + fileID: number; + score: number; +} + +export interface MLDataAPI { + // The whole stored ML payload for a file, or undefined when it is not + // cached. Reads the payload from disk, so it is async. + forFile(args: { fileID: number }): Promise; + // The files nearest the given file by cosine over their CLIP embeddings, + // most similar first, excluding the file itself. Empty when the file has + // no indexed embedding. + similar(args: { fileID: number; limit?: number }): SimilarResult[]; + // The files nearest a caller-supplied query embedding by cosine, most + // similar first. Empty when the query is the wrong length for the index, + // has zero magnitude, or the index is empty. + searchByEmbedding(args: { + embedding: ArrayLike; + limit?: number; + }): SimilarResult[]; +} + +// Rank the packed index by cosine similarity to `query`, most similar first, +// and return the top `limit`. `skip` (a query file's own id) is left out. Both +// each row's magnitude and the query's are computed here rather than cached: +// the index mutates as ML data is fetched, and one plain pass over ~50k×512 +// floats is fast enough that a norm cache would only add a staleness bug. A +// zero-magnitude vector has no direction, so it is dropped rather than divided +// by zero. +const topByCosine = ( + index: MLIndex, + query: ArrayLike, + limit: number, + skip?: number, +): SimilarResult[] => { + const { fileIDs, embeddingLength, embeddings } = index; + if (embeddingLength === 0 || query.length !== embeddingLength) return []; + + // Every indexed read below is in range: the inner loops run to + // `embeddingLength`, the query is exactly that long (checked above), and + // the packed buffer holds `fileIDs.length * embeddingLength` floats. + // `noUncheckedIndexedAccess` still widens each read to `number | undefined`, + // so they are asserted non-null rather than paying a per-element guard in + // this hot ~50k×512 loop. + let queryNorm = 0; + for (let k = 0; k < embeddingLength; k++) { + const q = query[k]!; + queryNorm += q * q; + } + queryNorm = Math.sqrt(queryNorm); + if (queryNorm === 0) return []; + + const results: SimilarResult[] = []; + for (let i = 0; i < fileIDs.length; i++) { + const id = fileIDs[i]!; + if (id === skip) continue; + const base = i * embeddingLength; + let dot = 0; + let norm = 0; + for (let k = 0; k < embeddingLength; k++) { + const v = embeddings[base + k]!; + dot += query[k]! * v; + norm += v * v; + } + if (norm === 0) continue; + results.push({ + fileID: id, + score: dot / (queryNorm * Math.sqrt(norm)), + }); + } + + // Descending score, ties broken by ascending fileID for a stable order. + results.sort((a, b) => b.score - a.score || a.fileID - b.fileID); + return results.slice(0, Math.max(0, Math.trunc(limit))); +}; + +// Build the search surface over a store the library supplies lazily (the store +// is absent when the client cannot fetch ML data). Reading it per call keeps +// the surface current as the index grows. +export const makeMLDataAPI = ( + store: () => MLDataStore | undefined, +): MLDataAPI => ({ + forFile: ({ fileID }): Promise => { + const s = store(); + return s ? s.readPayload(fileID) : Promise.resolve(undefined); + }, + similar: ({ fileID, limit }): SimilarResult[] => { + const s = store(); + if (!s) return []; + const index = s.getIndex(); + const pos = index.fileIDs.indexOf(fileID); + if (pos < 0) return []; + const base = pos * index.embeddingLength; + const query = index.embeddings.subarray( + base, + base + index.embeddingLength, + ); + return topByCosine(index, query, limit ?? DEFAULT_LIMIT, fileID); + }, + searchByEmbedding: ({ embedding, limit }): SimilarResult[] => { + const s = store(); + if (!s) return []; + return topByCosine(s.getIndex(), embedding, limit ?? DEFAULT_LIMIT); + }, +}); diff --git a/test/library/mlsearch.test.ts b/test/library/mlsearch.test.ts new file mode 100644 index 0000000..877fccc --- /dev/null +++ b/test/library/mlsearch.test.ts @@ -0,0 +1,108 @@ +/** + * Tests for the content-similarity search surface over the CLIP index + * (issue #50). + * + * The surface is `lib.mldata`: `forFile` reads the full stored payload from + * disk, while `similar` and `searchByEmbedding` rank fileIDs by cosine + * similarity over the in-RAM `Float32Array` index alone (no disk, no network). + * The fixture uses axis-aligned vectors so the correct cosine ranking is + * obvious by inspection; cosine ignores magnitude, so `[2, 0, 0]` ranks above + * `[0.8, 0.6, 0]` for a `[1, 0, 0]` query. + */ + +import { describe, it, expect, beforeEach, afterEach } from "vitest"; +import { mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; + +import { MLDataStore } from "../../src/library/mldata.js"; +import { makeMLDataAPI, type MLDataAPI } from "../../src/library/mlsearch.js"; +import type { MLData } from "../../src/mldata-fetch.js"; + +// A payload shaped like Ente's: a CLIP embedding plus face data that only the +// on-disk payload carries (never the RAM index). +const payload = (embedding: number[]): MLData => ({ + face: { faces: [{ faceID: "f", detection: { box: { x: 0.5 } } }] }, + clip: { embedding }, +}); + +// A small fixture index. Directions are chosen so every cosine ranking below +// is unambiguous. +const fixture = (): Map => + new Map([ + [10, payload([1, 0, 0])], + [20, payload([0.8, 0.6, 0])], + [30, payload([0, 1, 0])], + [40, payload([-1, 0, 0])], + [50, payload([2, 0, 0])], + ]); + +describe("lib.mldata content-similarity search", () => { + let dir: string; + let store: MLDataStore; + let api: MLDataAPI; + + beforeEach(async () => { + dir = mkdtempSync(join(tmpdir(), "quak-mlsearch-")); + store = await MLDataStore.open(dir); + const updation = new Map([...fixture().keys()].map((id) => [id, 1])); + await store.storeFetched(fixture(), updation); + api = makeMLDataAPI(() => store); + }); + + afterEach(() => { + rmSync(dir, { recursive: true, force: true }); + }); + + it("forFile returns the whole stored payload, or undefined when uncached", async () => { + const full = await api.forFile({ fileID: 20 }); + expect(full).toBeDefined(); + // Face data lives only in the payload, never in the RAM index. + expect(full?.face).toBeDefined(); + expect(full?.clip).toEqual({ embedding: [0.8, 0.6, 0] }); + expect(await api.forFile({ fileID: 999 })).toBeUndefined(); + }); + + it("similar ranks other files by cosine and excludes the query itself", () => { + // Query is file 10 = [1, 0, 0]. By cosine: 50 (1.0) > 20 (0.8) > + // 30 (0) > 40 (-1); 10 itself is left out. + const ranked = api.similar({ fileID: 10 }); + expect(ranked.map((r) => r.fileID)).toEqual([50, 20, 30, 40]); + // Cosine ignores magnitude: [2,0,0] is a perfect match for [1,0,0]. + expect(ranked[0]).toMatchObject({ fileID: 50 }); + expect(ranked[0].score).toBeCloseTo(1, 5); + }); + + it("similar honours limit and returns [] for an unindexed file", () => { + expect( + api.similar({ fileID: 10, limit: 2 }).map((r) => r.fileID), + ).toEqual([50, 20]); + expect(api.similar({ fileID: 999 })).toEqual([]); + }); + + it("searchByEmbedding ranks the index by cosine to the query vector", () => { + // Query [0, 1, 0]: 30 (1.0) > 20 (0.6) > {10, 40, 50} all 0, broken by + // ascending fileID. + const ranked = api.searchByEmbedding({ embedding: [0, 1, 0] }); + expect(ranked.map((r) => r.fileID)).toEqual([30, 20, 10, 40, 50]); + expect(ranked[0].score).toBeCloseTo(1, 5); + + expect( + api + .searchByEmbedding({ embedding: [0, 1, 0], limit: 2 }) + .map((r) => r.fileID), + ).toEqual([30, 20]); + }); + + it("searchByEmbedding returns [] for a wrong-length or zero query", () => { + expect(api.searchByEmbedding({ embedding: [1, 0] })).toEqual([]); + expect(api.searchByEmbedding({ embedding: [0, 0, 0] })).toEqual([]); + }); + + it("degrades to empty results when no ML store is present", async () => { + const none = makeMLDataAPI(() => undefined); + expect(await none.forFile({ fileID: 10 })).toBeUndefined(); + expect(none.similar({ fileID: 10 })).toEqual([]); + expect(none.searchByEmbedding({ embedding: [1, 0, 0] })).toEqual([]); + }); +});