quak/test/clip-keywords.test.ts

// `clipKeywords`: the keywords quak writes for a file in place of its CLIP
// embedding.
//
// Ente's ML data for a file holds a CLIP embedding, 512 numbers its image
// model gave the photo. Ente's photos app puts a file in a category of its
// Discover section when the cosine similarity of that embedding and the text
// embedding of the category's query is at least the category's minimum score,
// and keeps videos out of some categories. `clipKeywords` applies that rule
// and returns the titles of the categories a file falls in, highest score
// first; `quak backup` and `quak backup-metadata` write them as `keywords`.

import { readFileSync } from "node:fs";
import { describe, expect, it } from "vitest";

import { DISCOVER_CATEGORIES } from "../src/clip-keyword-vectors.js";
import { clipKeywords } from "../src/clip-keywords.js";

// 512 numbers: the ones given, then zeros.
const vector = (...start: number[]): number[] => [
    ...start,
    ...new Array<number>(512 - start.length).fill(0),
];

const withEmbedding = (embedding: unknown) => ({
    face: { faces: [] },
    clip: { embedding },
});

// Two made-up categories, each with a text embedding along one axis. The
// embedding [3, 4] has a cosine similarity of 3/5 with "Across" and 4/5 with
// "Up". Videos are kept out of "Across".
const ACROSS = {
    title: "Across",
    query: "across",
    minScore: 0.6,
    showVideo: false,
    textEmbedding: vector(1),
};
const UP = {
    title: "Up",
    query: "up",
    minScore: 0.5,
    showVideo: true,
    textEmbedding: vector(0, 1),
};

describe("clipKeywords", () => {
    it("gives each category whose minimum score the similarity reaches, highest score first", () => {
        expect(
            clipKeywords(withEmbedding(vector(3, 4)), false, [ACROSS, UP]),
        ).toEqual(["Up", "Across"]);
    });

    it("leaves out a category when the similarity is just below its minimum score", () => {
        const stricter = { ...ACROSS, minScore: 0.6 + 1e-9 };

        expect(
            clipKeywords(withEmbedding(vector(3, 4)), false, [stricter, UP]),
        ).toEqual(["Up"]);
    });

    it("keeps a video out of a category that leaves out videos", () => {
        expect(
            clipKeywords(withEmbedding(vector(3, 4)), true, [ACROSS, UP]),
        ).toEqual(["Up"]);
    });

    it("gives an empty list for an embedding that falls in no category", () => {
        expect(
            clipKeywords(withEmbedding(vector(0, 0, 1)), false, [ACROSS, UP]),
        ).toEqual([]);
    });

    it.each([
        ["no CLIP embedding", { face: { faces: [] } }],
        ["a CLIP part that is not an object", { clip: null }],
        ["a CLIP part without an embedding", { clip: {} }],
        ["an embedding of 511 numbers", withEmbedding(vector(3, 4).slice(1))],
        ["an embedding with NaN in it", withEmbedding(vector(3, NaN))],
        ["an embedding with Infinity in it", withEmbedding(vector(Infinity))],
        ["an embedding of strings", withEmbedding(vector(3, 4).map(String))],
        ["an embedding of zeros", withEmbedding(vector())],
    ])("gives no keywords for ML data with %s", (_, mlData) => {
        expect(clipKeywords(mlData, false, [ACROSS, UP])).toBeUndefined();
    });
});

// The CLIP embedding Ente's image model gives a public-domain NOAA photo of a
// sunset, made by bin/clip-keyword-vectors from the photo at `photo`.
const SUNSET: { photo: string; sha256: string; embedding: number[] } =
    JSON.parse(
        readFileSync(
            new URL("clip-keywords-sunset.json", import.meta.url),
            "utf-8",
        ),
    );

describe("Ente's Discover categories", () => {
    it("are the 16 of Ente's photos app, each with a text embedding of norm 1", () => {
        expect(DISCOVER_CATEGORIES.map((c) => c.title)).toEqual([
            "Identity",
            "Screenshots",
            "QR Codes",
            "Receipts",
            "Notes",
            "Memes",
            "Visiting Cards",
            "Babies",
            "Pets",
            "Selfies",
            "Wallpapers",
            "Food",
            "Celebrations",
            "Sunset",
            "Hills",
            "Greenery",
        ]);
        for (const category of DISCOVER_CATEGORIES) {
            expect(category.textEmbedding).toHaveLength(512);
            expect(Math.hypot(...category.textEmbedding)).toBeCloseTo(1, 5);
        }
    });

    it("give a photo of a sunset the keyword Sunset and no other", () => {
        expect(
            clipKeywords({ clip: { embedding: SUNSET.embedding } }, false),
        ).toEqual(["Sunset"]);
    });
});