sliftutils 1.4.81 → 1.4.83

package/storage/proxydatabase/inMemoryDatabase.ts

@@ -0,0 +1,108 @@
+import { Database } from "./Database";
+
+// Rough byte size of a value, for the read/written-bytes metrics.
+function sizeOf(value: unknown): number {
+    if (value === undefined || value === null) {
+        return 0;
+    }
+    if (value instanceof Uint8Array) {
+        return value.byteLength;
+    }
+    if (typeof value === "string") {
+        return value.length;
+    }
+    if (typeof value === "number") {
+        return 8;
+    }
+    if (typeof value === "boolean") {
+        return 1;
+    }
+    if (Array.isArray(value)) {
+        let total = 0;
+        for (const entry of value) {
+            total += sizeOf(entry);
+        }
+        return total;
+    }
+    if (typeof value === "object") {
+        let total = 0;
+        const record = value as Record<string, unknown>;
+        for (const key of Object.keys(record)) {
+            total += sizeOf(record[key]);
+        }
+        return total;
+    }
+    return 0;
+}
+
+// Record the property path a dereference lambda walks, so a write/delete can be applied at it.
+function capturePath(deref: (root: unknown) => unknown): string[] {
+    const path: string[] = [];
+    const handler: ProxyHandler<object> = {
+        get(_target, key) {
+            if (typeof key === "string") {
+                path.push(key);
+            }
+            return new Proxy({}, handler);
+        },
+    };
+    deref(new Proxy({}, handler));
+    return path;
+}
+
+function setPath(root: Record<string, unknown>, path: string[], value: unknown): void {
+    let node = root;
+    for (let index = 0; index < path.length - 1; index++) {
+        const key = path[index];
+        const next = node[key];
+        if (!next || typeof next !== "object") {
+            node[key] = {};
+        }
+        node = node[key] as Record<string, unknown>;
+    }
+    node[path[path.length - 1]] = value;
+}
+
+function deletePath(root: Record<string, unknown>, path: string[]): void {
+    let node = root;
+    for (let index = 0; index < path.length - 1; index++) {
+        const next = node[path[index]];
+        if (!next || typeof next !== "object") {
+            return;
+        }
+        node = next as Record<string, unknown>;
+    }
+    delete node[path[path.length - 1]];
+}
+
+// In-memory Database for tests/dev: holds the whole root as a plain object and counts every call plus the
+// bytes flowing through. Synchronous and always "synced", so reads only return undefined for a genuinely
+// missing scalar (a missing transaction set reads as empty inside the set helpers).
+export class InMemoryDatabase<Root> implements Database<Root> {
+    public readCalls = 0;
+    public writeCalls = 0;
+    public deleteCalls = 0;
+    public bytesRead = 0;
+    public bytesWritten = 0;
+    private root: Record<string, unknown>;
+
+    constructor(initial: Root) {
+        this.root = initial as Record<string, unknown>;
+    }
+
+    readData<Value>(deref: (root: Root) => Value): Value | undefined {
+        this.readCalls++;
+        const result = deref(this.root as Root);
+        this.bytesRead += sizeOf(result);
+        return result;
+    }
+    writeData<Value>(deref: (root: Root) => Value, newValue: Value): void {
+        this.writeCalls++;
+        this.bytesWritten += sizeOf(newValue);
+        setPath(this.root, capturePath(deref as (root: unknown) => unknown), newValue);
+    }
+    deleteData(deref: (root: Root) => unknown): void {
+        this.deleteCalls++;
+        deletePath(this.root, capturePath(deref as (root: unknown) => unknown));
+    }
+}

package/storage/proxydatabase/ivfDbCheck.d.ts

@@ -0,0 +1 @@
+export {};

package/storage/proxydatabase/ivfDbCheck.ts

@@ -0,0 +1,85 @@
+import fs from "fs";
+import { InMemoryDatabase } from "./inMemoryDatabase";
+import { IvfEmbeddingRoot, IvfConfig, EmbeddingInput, insertEmbeddings, searchEmbeddings } from "./ivfEmbeddingDatabase";
+import { encodeEmbedding, EmbeddingFormat } from "../embeddingFormats";
+import { ensureFaceData } from "../faceFramesData";
+
+// Exercises the tiered IVF embedding database through the in-memory Database wrapper, which counts every
+// read/write/delete call and the bytes through each, across dataset sizes and load patterns, with timing.
+
+const DIM = 512;
+const FILE_NAME = "faceFrames2_full.f32";
+const MODEL = "buffalo_l";
+const FORMAT: EmbeddingFormat = "q8g8_2048";
+
+async function loadFloats(count: number): Promise<Float32Array> {
+    const byteLength = count * DIM * 4;
+    const filePath = await ensureFaceData(FILE_NAME, byteLength);
+    const buffer = Buffer.alloc(byteLength);
+    const handle = fs.openSync(filePath, "r");
+    fs.readSync(handle, buffer, 0, byteLength, 0);
+    fs.closeSync(handle);
+    return new Float32Array(buffer.buffer, buffer.byteOffset, count * DIM);
+}
+
+async function buildInputs(count: number): Promise<EmbeddingInput[]> {
+    const floats = await loadFloats(count);
+    const inputs: EmbeddingInput[] = [];
+    for (let index = 0; index < count; index++) {
+        const slice = floats.subarray(index * DIM, (index + 1) * DIM);
+        inputs.push({ ref: "e" + index, embedding: encodeEmbedding({ input: slice, format: FORMAT, model: MODEL }) });
+    }
+    return inputs;
+}
+
+function freshDatabase(config: IvfConfig): InMemoryDatabase<IvfEmbeddingRoot> {
+    const root: IvfEmbeddingRoot = { config, count: 0, flat: {}, byRef: {}, steps: {}, centroids: {}, cells: {} };
+    return new InMemoryDatabase<IvfEmbeddingRoot>(root);
+}
+
+function pad(value: string | number, width: number): string {
+    return String(value).padStart(width);
+}
+
+async function main() {
+    const config: IvfConfig = { model: MODEL, format: FORMAT, cellTargetSize: 128 };
+    const runs: { size: number; batch: number }[] = [
+        { size: 10000, batch: 1 },
+        { size: 10000, batch: 100 },
+        { size: 100000, batch: 1000 },
+    ];
+
+    let maxSize = 0;
+    for (const run of runs) {
+        if (run.size > maxSize) {
+            maxSize = run.size;
+        }
+    }
+    const loadStart = Date.now();
+    const allInputs = await buildInputs(maxSize);
+    console.log(`loaded + encoded ${maxSize} embeddings in ${((Date.now() - loadStart) / 1000).toFixed(1)}s\n`);
+
+    console.log(`${pad("added", 7)} ${pad("batch", 5)} | ${pad("secs", 7)} ${pad("reads", 8)} ${pad("writes", 8)} ${pad("deletes", 8)} ${pad("readMB", 9)} ${pad("writeMB", 8)} | recall`);
+    for (const run of runs) {
+        const database = freshDatabase(config);
+        const start = Date.now();
+        for (let offset = 0; offset < run.size; offset += run.batch) {
+            insertEmbeddings(database, allInputs.slice(offset, offset + run.batch));
+        }
+        const seconds = (Date.now() - start) / 1000;
+
+        let found = 0;
+        const sampleCount = 100;
+        for (let index = 0; index < sampleCount; index++) {
+            const query = allInputs[Math.floor(index * run.size / sampleCount)];
+            const hits = searchEmbeddings(database, query.embedding, { probeBudget: 512, resultCount: 1 });
+            if (hits && hits.length && hits[0].ref === query.ref) {
+                found++;
+            }
+        }
+        const recall = (found / sampleCount).toFixed(2);
+        console.log(`${pad(run.size, 7)} ${pad(run.batch, 5)} | ${pad(seconds.toFixed(1), 7)} ${pad(database.readCalls, 8)} ${pad(database.writeCalls, 8)} ${pad(database.deleteCalls, 8)} ${pad((database.bytesRead / 1e6).toFixed(0), 9)} ${pad((database.bytesWritten / 1e6).toFixed(0), 8)} | ${recall}`);
+    }
+}
+
+void main();

package/storage/proxydatabase/ivfEmbeddingDatabase.d.ts

@@ -4,16 +4,22 @@ import { StoredEmbedding, EmbeddingFormat } from "../embeddingFormats";
 export type IvfConfig = {
     model: string;
     format: EmbeddingFormat;
-    cellTarget: number;
-    splitAt: number;
+    cellTargetSize: number;
 };
 export type IvfEmbeddingRoot = {
     config: IvfConfig;
-    centroids: TransactionSetStore;
+    count: number;
+    flat: TransactionSetStore<StoredEmbedding>;
+    byRef: {
+        [ref: string]: Uint8Array;
+    };
+    steps: {
+        [step: string]: boolean;
+    };
+    centroids: TransactionSetStore<StoredEmbedding>;
     cells: {
-        [cellId: string]: TransactionSetStore;
+        [cellId: string]: TransactionSetStore<StoredEmbedding>;
     };
-    refIndex: TransactionSetStore;
 };
 export type EmbeddingInput = {
     ref: string;
@@ -23,9 +29,11 @@ export type SearchHit = {
     ref: string;
     closeness: number;
 };
+export declare function rebuildStructure(database: Database<IvfEmbeddingRoot>): void;
 export declare function searchEmbeddings(database: Database<IvfEmbeddingRoot>, query: StoredEmbedding, options: {
     probeBudget: number;
     resultCount: number;
 }): SearchHit[] | undefined;
-export declare function insertEmbeddings(database: Database<IvfEmbeddingRoot>, items: EmbeddingInput[]): true | undefined;
-export declare function removeEmbeddings(database: Database<IvfEmbeddingRoot>, refs: string[]): true | undefined;
+export declare function lookupEmbeddings(database: Database<IvfEmbeddingRoot>, refs: string[]): Map<string, StoredEmbedding> | undefined;
+export declare function insertEmbeddings(database: Database<IvfEmbeddingRoot>, items: EmbeddingInput[]): undefined;
+export declare function removeEmbeddings(database: Database<IvfEmbeddingRoot>, refs: string[]): void;