A file's JSON holds keywords in place of its CLIP embedding (closes #218)

`quak backup` and `quak backup-metadata` write each file's ML data without its
CLIP embedding, and `keywords` in its place: the titles of the Discover
categories of Ente's photos app that the file falls in, highest score first. A
file falls in a category when the cosine similarity of its CLIP embedding and
the text embedding of the category's query reaches the category's minimum
score; as in Ente's app, five categories leave out videos.

The 16 text embeddings ship in `src/clip-keyword-vectors.ts`, written by
`bin/clip-keyword-vectors` from Ente's list and CLIP text model. quak runs no
model and gains no dependency.

Disclosure: the text embeddings are committed model output, against
REPO_POLICIES' rule on generated files.

Model: opus-5-5

test/cli/backup-end-to-end.test.ts

@@ -86,10 +86,17 @@ const DAY = "2024/2024-03/2024-03-01";
 const TOAST = new TextEncoder().encode("a few seconds of the wedding toast");
 
 // Ente's ML data for beach.jpg: a face it found, and the photo's CLIP
-// embedding.
+// embedding, here that of a photo of a sunset (see clip-keywords.test.ts).
 const ML_DATA = {
     face: { faces: [{ faceID: "100_1", score: 0.98 }] },
-    clip: { embedding: [0.5, 0.6, 0.7] },
+    clip: {
+        embedding: JSON.parse(
+            readFileSync(
+                new URL("../clip-keywords-sunset.json", import.meta.url),
+                "utf-8",
+            ),
+        ).embedding,
+    },
 };
 
 // Each file's ID, its owner, its original as uploaded, and the metadata its
@@ -486,7 +493,10 @@ describe("the first backup of the account", () => {
             },
             updationTime: TAKEN,
         });
-        expect(beach.mlData).toEqual(ML_DATA);
+        // The JSON holds the ML data without its CLIP embedding, and the
+        // keywords the embedding gives in its place.
+        expect(beach.mlData).toEqual({ face: ML_DATA.face });
+        expect(beach.keywords).toEqual(["Sunset"]);
         expect(beach.imageMetadata.exif.Make.description).toBe("Canon");
         expect(beach.imageMetadata.exif.Model.description).toBe("EOS R5");
 
@@ -497,6 +507,7 @@ describe("the first backup of the account", () => {
             metadata: { title: "toast.mp4", fileType: "video" },
         });
         expect(toast).not.toHaveProperty("mlData");
+        expect(toast).not.toHaveProperty("keywords");
         expect(toast).not.toHaveProperty("imageMetadata");
         expect(toast).not.toHaveProperty("imageMetadataError");
 

test/cli/backup.test.ts

@@ -52,6 +52,7 @@ import { tmpdir } from "node:os";
 import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
 
 import { runBackup, type BackupLibrary } from "../../src/backup.js";
+import { DISCOVER_CATEGORIES } from "../../src/clip-keyword-vectors.js";
 import { Library } from "../../src/library/index.js";
 import type { ContentSource } from "../../src/library/content.js";
 import type { CollectionsPage, FilesPage } from "../../src/client.js";
@@ -946,9 +947,19 @@ describe("the backup lock", () => {
     });
 });
 
-// Also serves ML data: file 100 has `ML_PAYLOAD`, the other two have none.
-// While `mlError` is set, every ML data request fails with it.
-const ML_PAYLOAD = { face: { faces: [] }, clip: { embedding: [0.5, 0.25] } };
+// Also serves ML data: file 100 has `ML_PAYLOAD`, whose CLIP embedding is that
+// of a photo of a sunset (see clip-keywords.test.ts), and the other two have
+// none. While `mlError` is set, every ML data request fails with it.
+const SUNSET_EMBEDDING: number[] = JSON.parse(
+    readFileSync(
+        new URL("../clip-keywords-sunset.json", import.meta.url),
+        "utf-8",
+    ),
+).embedding;
+const ML_PAYLOAD = {
+    face: { faces: [] },
+    clip: { embedding: SUNSET_EMBEDDING },
+};
 
 class MLClient extends MockClient {
     mlError?: string;
@@ -968,16 +979,99 @@ const fileJSON = (outDir: string, fileID: number): Record<string, unknown> =>
     JSON.parse(readFileSync(saved(outDir, `${fileID}.json`), "utf-8"));
 
 describe("ML data in each file's JSON", () => {
-    it("writes a file's ML data, and no ML field for a file that has none", async () => {
+    // The JSON holds the ML data without its CLIP embedding, and in its place
+    // the keywords the embedding gives.
+    it("writes a file's ML data and keywords, and no ML field for a file that has none", async () => {
         const lib = await openLibrary(stubSource(), new MLClient());
         const outDir = join(root, "backup");
 
         const result = await lib.backup({ downloadDirectory: outDir });
 
         expect(result.failed).toBe(0);
-        expect(fileJSON(outDir, 100).mlData).toEqual(ML_PAYLOAD);
+        expect(fileJSON(outDir, 100).mlData).toEqual({ face: { faces: [] } });
+        expect(fileJSON(outDir, 100).keywords).toEqual(["Sunset"]);
         expect(fileJSON(outDir, 101)).not.toHaveProperty("mlData");
         expect(fileJSON(outDir, 101)).not.toHaveProperty("mlDataError");
+        expect(fileJSON(outDir, 101)).not.toHaveProperty("keywords");
+        await lib.close();
+    });
+
+    // A file's ML data can hold faces and no CLIP embedding. Its JSON then has
+    // the ML data and no keywords.
+    it("writes no keywords for ML data without a CLIP embedding", async () => {
+        const faces = { face: { faces: [{ faceID: "101_1", score: 0.9 }] } };
+
+        // Gives file 101 ML data holding only faces.
+        class FacesClient extends MockClient {
+            async fetchMLData(args: { fileIDs: number[] }): Promise<{
+                payloads: Map<number, MLData>;
+                errors: Map<number, unknown>;
+            }> {
+                const payloads = new Map<number, MLData>();
+                if (args.fileIDs.includes(101)) payloads.set(101, faces);
+                return { payloads, errors: new Map() };
+            }
+        }
+        const lib = await openLibrary(stubSource(), new FacesClient());
+        const outDir = join(root, "backup");
+
+        const result = await lib.backup({ downloadDirectory: outDir });
+
+        expect(result.failed).toBe(0);
+        expect(fileJSON(outDir, 101).mlData).toEqual(faces);
+        expect(fileJSON(outDir, 101)).not.toHaveProperty("keywords");
+        await lib.close();
+    });
+
+    // Ente's photos app keeps videos out of some Discover categories, Identity
+    // among them. Here a photo and a video share one CLIP embedding, Identity's
+    // own text embedding, so the photo gets the keyword Identity and the video
+    // does not.
+    it("keeps a video out of a category that leaves out videos", async () => {
+        const identity = DISCOVER_CATEGORIES.find(
+            (c) => c.title === "Identity",
+        )!;
+        expect(identity.showVideo).toBe(false);
+        const toast = file(102, 1, "toast.mp4");
+        const video = {
+            ...toast,
+            metadata: { ...toast.metadata, fileType: "video" as const },
+        };
+
+        // Serves beach.jpg and the video in the Vacation album, and gives
+        // each ML data holding only that embedding.
+        class VideoClient extends MockClient {
+            override async filesSince(args: {
+                collectionID: number;
+            }): Promise<FilesPage> {
+                const files =
+                    args.collectionID === 1
+                        ? [file(100, 1, "beach.jpg"), video]
+                        : [];
+                return { files, deleted: [], cursor: 1 };
+            }
+            async fetchMLData(args: { fileIDs: number[] }): Promise<{
+                payloads: Map<number, MLData>;
+                errors: Map<number, unknown>;
+            }> {
+                const payloads = new Map<number, MLData>();
+                for (const id of args.fileIDs) {
+                    payloads.set(id, {
+                        clip: { embedding: identity.textEmbedding },
+                    });
+                }
+                return { payloads, errors: new Map() };
+            }
+        }
+        const lib = await openLibrary(stubSource(), new VideoClient());
+        const outDir = join(root, "backup");
+
+        const result = await lib.backup({ downloadDirectory: outDir });
+
+        expect(result.failed).toBe(0);
+        expect(fileJSON(outDir, 100).keywords).toContain("Identity");
+        expect(fileJSON(outDir, 102)).toHaveProperty("keywords");
+        expect(fileJSON(outDir, 102).keywords).not.toContain("Identity");
         await lib.close();
     });
 
@@ -993,7 +1087,7 @@ describe("ML data in each file's JSON", () => {
         const result = await lib.backup({ downloadDirectory: outDir });
 
         // File 100's ML data was cached by the first run and is kept.
-        expect(fileJSON(outDir, 100).mlData).toEqual(ML_PAYLOAD);
+        expect(fileJSON(outDir, 100).mlData).toEqual({ face: { faces: [] } });
         expect(fileJSON(outDir, 100)).not.toHaveProperty("mlDataError");
         for (const fileID of [101, 200]) {
             expect(fileJSON(outDir, fileID).mlDataError).toBe(
@@ -1026,7 +1120,7 @@ describe("ML data in each file's JSON", () => {
 
         expect(result.failed).toBe(0);
         expect(result.skipped).toBe(3);
-        expect(fileJSON(outDir, 100).mlData).toEqual(ML_PAYLOAD);
+        expect(fileJSON(outDir, 100).mlData).toEqual({ face: { faces: [] } });
         for (const fileID of [100, 101, 200]) {
             expect(fileJSON(outDir, fileID)).not.toHaveProperty("mlDataError");
         }

test/cli/metadata-backup.test.ts

@@ -41,6 +41,7 @@ import { beforeAll, afterAll, describe, expect, it, vi } from "vitest";
 import { init, toBase64, encryptBlob } from "../../src/crypto/index.js";
 import * as jpegJs from "jpeg-js";
 import { Client } from "../../src/client.js";
+import { DISCOVER_CATEGORIES } from "../../src/clip-keyword-vectors.js";
 import { Library } from "../../src/library/index.js";
 import {
     runMetadataBackup,
@@ -56,8 +57,8 @@ import {
     type MetaMockState,
 } from "../ente-server.js";
 
-// One file per ML data request, so the two files of the mock account are
-// fetched in two requests and one of them can fail on its own.
+// One file per ML data request, so each file of the mock account is fetched in
+// a request of its own and one of them can fail on its own.
 vi.mock("../../src/mldata-fetch.js", async (importOriginal) => ({
     ...(await importOriginal<typeof import("../../src/mldata-fetch.js")>()),
     MLDATA_BATCH_SIZE: 1,
@@ -66,6 +67,21 @@ vi.mock("../../src/mldata-fetch.js", async (importOriginal) => ({
 const TEST_EMAIL = "metabackup@example.com";
 const TEST_PASSWORD = "metapass";
 
+// The CLIP embedding of a photo of a sunset (see clip-keywords.test.ts), which
+// the ML data of beach.jpg holds.
+const SUNSET_EMBEDDING: number[] = JSON.parse(
+    readFileSync(
+        new URL("../clip-keywords-sunset.json", import.meta.url),
+        "utf-8",
+    ),
+).embedding;
+
+// Identity, a Discover category of Ente's photos app that leaves out videos.
+const IDENTITY = DISCOVER_CATEGORIES.find((c) => c.title === "Identity")!;
+
+// ML data holding faces and no CLIP embedding, which file 302 has.
+const FACES_ONLY = { face: { faces: [{ faceID: "face-def", score: 0.9 }] } };
+
 let mock: MetaMockState;
 let testDir: string;
 
@@ -204,6 +220,65 @@ const buildMetaMock = async (): Promise<MetaMockState> => {
         updationTime: 1710000000000000,
     };
 
+    // A file in coll 10 with the ML data given. `fileType` is Ente's number: 0
+    // an image, 1 a video.
+    const vacationFile = (
+        id: number,
+        title: string,
+        fileType: number,
+        mlData: object,
+    ) => {
+        const fileKey = sodium.crypto_secretstream_xchacha20poly1305_keygen();
+        const { ciphertext: encKey, nonce } = encryptSecretbox(fileKey, ck1);
+        const meta = encryptStreamBlob(
+            new TextEncoder().encode(
+                JSON.stringify({
+                    title,
+                    fileType,
+                    creationTime: 1700000000000000,
+                    modificationTime: 1700000000000000,
+                }),
+            ),
+            fileKey,
+        );
+        const ml = encryptBlob(
+            new Uint8Array(gzipSync(Buffer.from(JSON.stringify(mlData)))),
+            fileKey,
+        );
+        return {
+            raw: {
+                id,
+                collectionID: 10,
+                ownerID: 42,
+                encryptedKey: toBase64(encKey),
+                keyDecryptionNonce: toBase64(nonce),
+                metadata: {
+                    encryptedData: toBase64(meta.ciphertext),
+                    decryptionHeader: toBase64(meta.header),
+                },
+                file: {
+                    decryptionHeader: toBase64(sodium.randombytes_buf(24)),
+                },
+                thumbnail: {
+                    decryptionHeader: toBase64(sodium.randombytes_buf(24)),
+                },
+                updationTime: 1700000000000000,
+            },
+            mlData: {
+                encryptedData: toBase64(ml.ciphertext),
+                decryptionHeader: toBase64(ml.header),
+            },
+        };
+    };
+    // Files 300 and 301 in coll 10: a photo and a video whose ML data holds
+    // only a CLIP embedding, Identity's own text embedding.
+    const identityML = { clip: { embedding: IDENTITY.textEmbedding } };
+    const identityPhoto = vacationFile(300, "passport.jpg", 0, identityML);
+    const identityVideo = vacationFile(301, "passport.mp4", 1, identityML);
+    // File 302 in coll 10: a photo whose ML data holds faces and no CLIP
+    // embedding.
+    const facesPhoto = vacationFile(302, "party.jpg", 0, FACES_ONLY);
+
     // Encrypt ML data for file 100 (gzipped JSON, encrypted with file key)
     const mlPayload = JSON.stringify({
         face: {
@@ -230,7 +305,7 @@ const buildMetaMock = async (): Promise<MetaMockState> => {
         clip: {
             version: 1,
             client: "test",
-            embedding: [0.5, 0.6, 0.7],
+            embedding: SUNSET_EMBEDDING,
         },
     });
     const gzipped = gzipSync(Buffer.from(mlPayload));
@@ -270,12 +345,23 @@ const buildMetaMock = async (): Promise<MetaMockState> => {
         keyAttributes: account.keyAttributes,
         encryptedToken: account.encryptedToken,
         collections: [rawColl1, rawColl2],
-        filesByCollection: { 10: [rawFile1], 20: [rawFile2] },
+        filesByCollection: {
+            10: [
+                rawFile1,
+                identityPhoto.raw,
+                identityVideo.raw,
+                facesPhoto.raw,
+            ],
+            20: [rawFile2],
+        },
         encryptedMLData: {
             100: {
                 encryptedData: toBase64(mlCiphertext),
                 decryptionHeader: toBase64(mlHeader),
             },
+            300: identityPhoto.mlData,
+            301: identityVideo.mlData,
+            302: facesPhoto.mlData,
         },
         fileCiphertexts: { 100: encFileBody1 },
     };
@@ -432,7 +518,9 @@ describe("quak backup-metadata", () => {
         expect(account.email).toBe(TEST_EMAIL);
     });
 
-    it("fetches and decrypts ML data by default", async () => {
+    // The CLIP embedding is left out of the ML data, and the keywords it gives
+    // are written in its place.
+    it("fetches and decrypts ML data by default, with keywords for the CLIP embedding", async () => {
         const outDir = join(testDir, "ml-data");
         await runBackup(outDir);
 
@@ -452,8 +540,41 @@ describe("quak backup-metadata", () => {
         expect(fileMeta.mlData.face.faces[0].faceID).toBe("face-abc");
         expect(fileMeta.mlData.face.faces[0].score).toBeCloseTo(0.98);
         expect(fileMeta.mlData.face.faces[0].detection.box.x).toBeCloseTo(0.1);
-        expect(fileMeta.mlData.clip).toBeDefined();
-        expect(fileMeta.mlData.clip.embedding).toEqual([0.5, 0.6, 0.7]);
+        expect(fileMeta.mlData).not.toHaveProperty("clip");
+        expect(fileMeta.keywords).toEqual(["Sunset"]);
+    });
+
+    // Ente's photos app keeps videos out of some Discover categories, Identity
+    // among them. Files 300 and 301, a photo and a video, share one CLIP
+    // embedding, Identity's own text embedding, so the photo gets the keyword
+    // Identity and the video does not.
+    it("keeps a video out of a category that leaves out videos", async () => {
+        expect(IDENTITY.showVideo).toBe(false);
+        const outDir = join(testDir, "video-keywords");
+        await runBackup(outDir);
+
+        const vacation = join(outDir, "collections", "10-Vacation");
+        const fileJSON = (fileID: number) =>
+            JSON.parse(readFileSync(join(vacation, `${fileID}.json`), "utf-8"));
+        expect(fileJSON(300).keywords).toContain("Identity");
+        expect(fileJSON(301)).toHaveProperty("keywords");
+        expect(fileJSON(301).keywords).not.toContain("Identity");
+    });
+
+    // A file's ML data can hold faces and no CLIP embedding. File 302's JSON
+    // then has the ML data and no keywords.
+    it("writes no keywords for ML data without a CLIP embedding", async () => {
+        const outDir = join(testDir, "no-clip");
+        await runBackup(outDir);
+
+        const fileMeta = JSON.parse(
+            readFileSync(
+                join(outDir, "collections", "10-Vacation", "302.json"),
+                "utf-8",
+            ),
+        );
+        expect(fileMeta.mlData).toEqual(FACES_ONLY);
+        expect(fileMeta).not.toHaveProperty("keywords");
     });
 
     it("extracts EXIF from downloaded files when --exif is set", async () => {
@@ -541,7 +662,7 @@ describe("quak backup-metadata when a file's ML data cannot be read", () => {
                 "utf-8",
             ),
         );
-        expect(ok.mlData.clip.embedding).toEqual([0.5, 0.6, 0.7]);
+        expect(ok.keywords).toEqual(["Sunset"]);
         expect(ok.mlDataError).toBeUndefined();
 
         const failed = JSON.parse(
@@ -599,7 +720,7 @@ describe("quak backup-metadata when a file's ML data cannot be read", () => {
                 "utf-8",
             ),
         );
-        expect(ok.mlData.clip.embedding).toEqual([0.5, 0.6, 0.7]);
+        expect(ok.keywords).toEqual(["Sunset"]);
     });
 
     it("exits 0 when every ML data request succeeds", async () => {