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

src/backup.ts

@@ -14,9 +14,9 @@
 //         YYYY/YYYY-MM/YYYY-MM-DD/
 //             YYYY-MM-DD.<fileID>.<ext>    the decrypted bytes (the save path)
 //             YYYY-MM-DD.<fileID>.json     per-file metadata sidecar, with
-//                                          the file's ML data and its
-//                                          original's EXIF, XMP and
-//                                          dimensions
+//                                          the file's ML data, its
+//                                          keywords, and its original's
+//                                          EXIF, XMP and dimensions
 //         collections/<name>/<title>       symlink to the original
 //         collections/<name>.json          per-collection metadata
 //         account.json                     the account's email and user ID
@@ -77,6 +77,7 @@ import { readFile } from "node:fs/promises";
 import { dirname, extname, join, relative, resolve } from "node:path";
 import lockfile from "proper-lockfile";
 
+import { clipKeywords } from "./clip-keywords.js";
 import { removeLeftoverTempFiles } from "./download/index.js";
 import { describeError } from "./errors.js";
 import {
@@ -497,8 +498,9 @@ const imageMetadataFor = async (
     }
 };
 
-// The file's JSON: its basic fields, its magic metadata, its ML data or the
-// reason the ML data is missing, and its image metadata.
+// The file's JSON: its basic fields, its magic metadata, its ML data without
+// the CLIP embedding and the keywords that embedding gives, or the reason the
+// ML data is missing, and its image metadata.
 const writeSidecar = (
     path: string,
     file: EnteFile,
@@ -514,7 +516,13 @@ const writeSidecar = (
     };
     if (file.magicMetadata) meta.magicMetadata = file.magicMetadata;
     if (file.pubMagicMetadata) meta.pubMagicMetadata = file.pubMagicMetadata;
-    if (ml.mlData) meta.mlData = ml.mlData;
+    if (ml.mlData) {
+        const { clip: _clip, ...mlDataWithoutClip } = ml.mlData;
+        meta.mlData = mlDataWithoutClip;
+        const isVideo = file.metadata.fileType === "video";
+        const keywords = clipKeywords(ml.mlData, isVideo);
+        if (keywords) meta.keywords = keywords;
+    }
     if (ml.mlDataError) meta.mlDataError = ml.mlDataError;
     if (image.imageMetadata) meta.imageMetadata = image.imageMetadata;
     if (image.imageMetadataError) {

src/clip-keywords.ts

@@ -0,0 +1,57 @@
+// Keywords for a file from its CLIP embedding.
+//
+// Ente stores no keywords. A file's ML data holds a CLIP embedding, the 512
+// numbers Ente's image model gave the photo, and Ente's photos app sorts files
+// into the categories of its Discover section with it: a file falls in a
+// category when the cosine similarity of its CLIP embedding and the text
+// embedding of the category's query is at least the category's `minScore`,
+// unless the file is a video and the category leaves out videos. A file's
+// keywords are the titles of the categories it falls in. The categories and
+// their text embeddings are in `clip-keyword-vectors.ts`.
+
+import { DISCOVER_CATEGORIES } from "./clip-keyword-vectors.js";
+import { clipEmbedding } from "./library/mldata.js";
+import type { MLData } from "./mldata-fetch.js";
+
+// How many numbers a CLIP embedding from Ente's image model has.
+const EMBEDDING_LENGTH = 512;
+
+const norm = (vector: number[]): number => Math.hypot(...vector);
+
+const cosineSimilarity = (a: number[], b: number[]): number => {
+    let dot = 0;
+    a.forEach((value, i) => {
+        dot += value * (b[i] ?? 0);
+    });
+    return dot / (norm(a) * norm(b));
+};
+
+// The titles of the Discover categories a file falls in, highest score first,
+// or undefined when its ML data has no CLIP embedding of 512 finite numbers
+// that are not all zero. Tests pass their own `categories`.
+export const clipKeywords = (
+    mlData: MLData,
+    isVideo: boolean,
+    categories = DISCOVER_CATEGORIES,
+): string[] | undefined => {
+    const embedding = clipEmbedding(mlData);
+    if (
+        !embedding ||
+        embedding.length !== EMBEDDING_LENGTH ||
+        norm(embedding) === 0
+    ) {
+        return undefined;
+    }
+
+    const matches: { title: string; score: number }[] = [];
+    for (const category of categories) {
+        if (isVideo && !category.showVideo) continue;
+        const score = cosineSimilarity(embedding, category.textEmbedding);
+        if (score >= category.minScore) {
+            matches.push({ title: category.title, score });
+        }
+    }
+    return matches
+        .sort((a, b) => b.score - a.score)
+        .map((match) => match.title);
+};

src/library/mldata.ts

@@ -61,7 +61,7 @@ export interface MLIndex {
 // Pull the CLIP embedding out of a payload, or undefined when it is absent or
 // misshapen. Kept strict so a bad payload is skipped rather than corrupting the
 // packed index.
-const clipEmbedding = (payload: MLData): number[] | undefined => {
+export const clipEmbedding = (payload: MLData): number[] | undefined => {
     const clip = payload.clip;
     if (typeof clip !== "object" || clip === null) return undefined;
     const embedding = (clip as { embedding?: unknown }).embedding;

src/metadata-backup.ts

@@ -1,6 +1,7 @@
 import { mkdirSync, readFileSync, writeFileSync } from "node:fs";
 import { join } from "node:path";
 import * as jpeg from "jpeg-js";
+import { clipKeywords } from "./clip-keywords.js";
 import type { Client } from "./client.js";
 import { describeError } from "./errors.js";
 import { readExifTags } from "./exif.js";
@@ -97,7 +98,8 @@ const extractExif = async (
 
 // Dump every decrypted metadata layer the account holds into a directory tree
 // of plain JSON: account, per-collection, and per-file records including the
-// private and public magic metadata and (by default) the ML data. Collections
+// private and public magic metadata and the ML data, with the keywords of its
+// CLIP embedding in place of the embedding (see clip-keywords.ts). Collections
 // and files are enumerated from the library's cache, which the caller refreshes
 // first. Returns how many files have no ML data because their ML data request
 // failed or their ML data did not decrypt, gunzip or parse; they are still
@@ -230,8 +232,16 @@ export const runMetadataBackup = async (
         if (file.pubMagicMetadata)
             fileMeta.pubMagicMetadata = file.pubMagicMetadata;
 
+        // The ML data without the CLIP embedding, and the keywords that
+        // embedding gives.
         const ml = mlDataMap.get(file.id);
-        if (ml) fileMeta.mlData = ml;
+        if (ml) {
+            const { clip: _clip, ...mlDataWithoutClip } = ml;
+            fileMeta.mlData = mlDataWithoutClip;
+            const isVideo = file.metadata.fileType === "video";
+            const keywords = clipKeywords(ml, isVideo);
+            if (keywords) fileMeta.keywords = keywords;
+        }
         const mlError = mlDataErrors.get(file.id);
         if (mlError) fileMeta.mlDataError = mlError;