diff --git a/src/library/index.ts b/src/library/index.ts index ba2a315..05a3965 100644 --- a/src/library/index.ts +++ b/src/library/index.ts @@ -41,6 +41,7 @@ import { type PhotosAPI, type TimelineAPI, } from "./read.js"; +import { makeMLDataAPI, type MLDataAPI } from "./mlsearch.js"; export { Album, @@ -52,6 +53,7 @@ export { type TimelineGroup, type GroupBy, } from "./read.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"; @@ -138,6 +140,10 @@ export class Library { readonly albums: AlbumsAPI; readonly photos: PhotosAPI; readonly timeline: TimelineAPI; + // 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; @@ -146,7 +152,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; @@ -188,7 +194,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.lastRecords = this.deriveNow(); // The read namespaces derive fresh from the store on each call, so they @@ -197,6 +203,8 @@ export class Library { this.albums = makeAlbumsAPI(derive); this.photos = makePhotosAPI(derive); this.timeline = makeTimelineAPI(derive); + // 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 @@ -293,7 +301,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(); return { userID: this.store.userID, collections: collections.length, @@ -448,7 +456,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..cdfeae0 --- /dev/null +++ b/src/library/mlsearch.ts @@ -0,0 +1,120 @@ +// 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 []; + + let queryNorm = 0; + for (let k = 0; k < embeddingLength; k++) queryNorm += query[k] * query[k]; + 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([]); + }); +});