Content-similarity search surface over the CLIP index (closes #50)
check / check (push) Failing after 25s
check / check (push) Failing after 25s
Add `lib.mldata`, the RAM-only content search over the #49 CLIP index: - `forFile({ fileID })` — the whole stored ML payload from disk, or undefined when uncached (the only method here that touches disk). - `similar({ fileID, limit? })` — files nearest a file by cosine over their CLIP embeddings, excluding the file itself. - `searchByEmbedding({ embedding, limit? })` — files nearest a caller-supplied query vector; quak bundles no text encoder (owner-deferred). Cosine is one plain loop over the packed Float32Array index, computing each row's magnitude on the fly rather than caching norms the mutating index would invalidate. No new dependency. Results are `{ fileID, score }[]` descending, ties broken by ascending fileID. The private store field was renamed `mlStore` to free `mldata` for the public API. Model: opus-4-8
This commit is contained in:
+12
-4
@@ -41,6 +41,7 @@ import {
|
|||||||
type PhotosAPI,
|
type PhotosAPI,
|
||||||
type TimelineAPI,
|
type TimelineAPI,
|
||||||
} from "./read.js";
|
} from "./read.js";
|
||||||
|
import { makeMLDataAPI, type MLDataAPI } from "./mlsearch.js";
|
||||||
|
|
||||||
export {
|
export {
|
||||||
Album,
|
Album,
|
||||||
@@ -52,6 +53,7 @@ export {
|
|||||||
type TimelineGroup,
|
type TimelineGroup,
|
||||||
type GroupBy,
|
type GroupBy,
|
||||||
} from "./read.js";
|
} from "./read.js";
|
||||||
|
export { type MLDataAPI, type SimilarResult } from "./mlsearch.js";
|
||||||
import type { CollectionsPage, FilesPage } from "../client.js";
|
import type { CollectionsPage, FilesPage } from "../client.js";
|
||||||
import { MLDATA_BATCH_SIZE, type MLData } from "../mldata-fetch.js";
|
import { MLDATA_BATCH_SIZE, type MLData } from "../mldata-fetch.js";
|
||||||
import type { Collection, EnteFile } from "../model/types.js";
|
import type { Collection, EnteFile } from "../model/types.js";
|
||||||
@@ -138,6 +140,10 @@ export class Library {
|
|||||||
readonly albums: AlbumsAPI;
|
readonly albums: AlbumsAPI;
|
||||||
readonly photos: PhotosAPI;
|
readonly photos: PhotosAPI;
|
||||||
readonly timeline: TimelineAPI;
|
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 client: LibraryClient;
|
||||||
private readonly store: MetadataStore;
|
private readonly store: MetadataStore;
|
||||||
@@ -146,7 +152,7 @@ export class Library {
|
|||||||
private readonly onProgress?: RefreshProgressCallback;
|
private readonly onProgress?: RefreshProgressCallback;
|
||||||
private readonly pools: RequestPools;
|
private readonly pools: RequestPools;
|
||||||
// The ML-data cache, present only when the client can fetch ML data.
|
// The ML-data cache, present only when the client can fetch ML data.
|
||||||
private readonly mldata?: MLDataStore;
|
private readonly mlStore?: MLDataStore;
|
||||||
|
|
||||||
private timer?: ReturnType<typeof setTimeout>;
|
private timer?: ReturnType<typeof setTimeout>;
|
||||||
private refreshing = false;
|
private refreshing = false;
|
||||||
@@ -188,7 +194,7 @@ export class Library {
|
|||||||
this.intervalMs = args.intervalMs;
|
this.intervalMs = args.intervalMs;
|
||||||
this.onProgress = args.onProgress;
|
this.onProgress = args.onProgress;
|
||||||
this.pools = args.pools;
|
this.pools = args.pools;
|
||||||
this.mldata = args.mldata;
|
this.mlStore = args.mldata;
|
||||||
this.lastRecords = this.deriveNow();
|
this.lastRecords = this.deriveNow();
|
||||||
|
|
||||||
// The read namespaces derive fresh from the store on each call, so they
|
// 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.albums = makeAlbumsAPI(derive);
|
||||||
this.photos = makePhotosAPI(derive);
|
this.photos = makePhotosAPI(derive);
|
||||||
this.timeline = makeTimelineAPI(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
|
// 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) {
|
for (const c of collections) {
|
||||||
files += this.store.listFiles(c.id).length;
|
files += this.store.listFiles(c.id).length;
|
||||||
}
|
}
|
||||||
const ml = this.mldata?.stats();
|
const ml = this.mlStore?.stats();
|
||||||
return {
|
return {
|
||||||
userID: this.store.userID,
|
userID: this.store.userID,
|
||||||
collections: collections.length,
|
collections: collections.length,
|
||||||
@@ -448,7 +456,7 @@ export class Library {
|
|||||||
// advanced), through the metadata pool, and update the CLIP index. Guarded
|
// advanced), through the metadata pool, and update the CLIP index. Guarded
|
||||||
// so passes never overlap; a failure is reported, not thrown.
|
// so passes never overlap; a failure is reported, not thrown.
|
||||||
private async runMLFetch(): Promise<void> {
|
private async runMLFetch(): Promise<void> {
|
||||||
const mldata = this.mldata;
|
const mldata = this.mlStore;
|
||||||
// Bind so the call keeps the client as its receiver when invoked
|
// Bind so the call keeps the client as its receiver when invoked
|
||||||
// through the pool below.
|
// through the pool below.
|
||||||
const fetchMLData = this.client.fetchMLData?.bind(this.client);
|
const fetchMLData = this.client.fetchMLData?.bind(this.client);
|
||||||
|
|||||||
@@ -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<MLData | undefined>;
|
||||||
|
// 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<number>;
|
||||||
|
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<number>,
|
||||||
|
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<MLData | undefined> => {
|
||||||
|
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);
|
||||||
|
},
|
||||||
|
});
|
||||||
@@ -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<number, MLData> =>
|
||||||
|
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([]);
|
||||||
|
});
|
||||||
|
});
|
||||||
Reference in New Issue
Block a user