toy_crossq.py

"""Toy proxy for CrossQ's late-interaction conditional quantization idea.

This is deliberately not a reproduction of the paper's benchmark.  It uses
synthetic topic-conditioned token embeddings to test the mechanism and report
the resulting index-size trade-off.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np


def normalize(x: np.ndarray) -> np.ndarray:
    return x / np.maximum(np.linalg.norm(x, axis=-1, keepdims=True), 1e-8)


def maxsim(queries: np.ndarray, docs: np.ndarray) -> np.ndarray:
    q = normalize(queries)
    d = normalize(docs)
    token_scores = np.einsum("qad,nbd->qnab", q, d)
    return np.max(np.max(token_scores, axis=3), axis=2)


def quantize(values: np.ndarray, low: np.ndarray, high: np.ndarray, bits: int) -> np.ndarray:
    levels = (1 << bits) - 1
    scaled = np.clip((values - low) / np.maximum(high - low, 1e-8), 0, 1)
    return np.rint(scaled * levels).astype(np.uint8)


def dequantize(codes: np.ndarray, low: np.ndarray, high: np.ndarray, bits: int) -> np.ndarray:
    levels = (1 << bits) - 1
    return low + codes.astype(np.float32) / levels * (high - low)


def reciprocal_rank(scores: np.ndarray, topics: np.ndarray, query_topics: np.ndarray) -> float:
    ranks = []
    for row, topic in zip(scores, query_topics):
        order = np.argsort(-row)
        relevant = np.flatnonzero(topics[order] == topic)
        ranks.append(1.0 / (int(np.flatnonzero(relevant)[0]) + 1) if relevant.size else 0.0)
    return float(np.mean(ranks))


def run(seed: int, documents: int, queries: int, bits: int) -> dict[str, float | int | str]:
    rng = np.random.default_rng(seed)
    tokens, dimensions, topic_count = 8, 32, 12
    centroids = normalize(rng.normal(size=(topic_count, dimensions)).astype(np.float32))
    doc_topics = rng.integers(topic_count, size=documents)
    query_topics = rng.integers(topic_count, size=queries)
    docs = centroids[doc_topics, None, :] + 0.20 * rng.normal(size=(documents, tokens, dimensions))
    qs = centroids[query_topics, None, :] + 0.20 * rng.normal(size=(queries, 4, dimensions))
    docs = docs.astype(np.float32)
    qs = qs.astype(np.float32)

    full_scores = maxsim(qs, docs)
    full_mrr = reciprocal_rank(full_scores, doc_topics, query_topics)

    global_low, global_high = docs.min(axis=(0, 1)), docs.max(axis=(0, 1))
    global_codes = quantize(docs, global_low, global_high, bits)
    global_docs = dequantize(global_codes, global_low, global_high, bits)
    global_mrr = reciprocal_rank(maxsim(qs, global_docs), doc_topics, query_topics)

    # Context is the document mean. Residual token codes are conditioned on it.
    context = docs.mean(axis=1)
    residual = docs - context[:, None, :]
    residual_low = residual.min(axis=(1, 2), keepdims=True)
    residual_high = residual.max(axis=(1, 2), keepdims=True)
    residual_codes = quantize(residual, residual_low, residual_high, bits)
    conditional_docs = context[:, None, :] + dequantize(
        residual_codes, residual_low, residual_high, bits
    )
    conditional_mrr = reciprocal_rank(
        maxsim(qs, conditional_docs), doc_topics, query_topics
    )

    full_bits = tokens * dimensions * 32
    conditional_bits = tokens * dimensions * bits + dimensions * 8 + 2 * 32
    compression = full_bits / conditional_bits
    return {
        "status": "toy",
        "seed": seed,
        "documents": documents,
        "queries": queries,
        "tokens_per_document": tokens,
        "dimensions": dimensions,
        "bits": bits,
        "full_precision_mrr10_proxy": round(full_mrr, 6),
        "global_quantization_mrr10_proxy": round(global_mrr, 6),
        "conditional_quantization_mrr10_proxy": round(conditional_mrr, 6),
        "conditional_minus_global": round(conditional_mrr - global_mrr, 6),
        "conditional_index_compression_x": round(compression, 3),
        "paper_claim_61x_verified": False,
        "paper_claim_2_3_percent_verified": False,
        "note": "Synthetic proxy only; paper code, weights, and benchmark data were unavailable.",
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--seed", type=int, default=4231)
    parser.add_argument("--documents", type=int, default=120)
    parser.add_argument("--queries", type=int, default=48)
    parser.add_argument("--bits", type=int, default=2)
    parser.add_argument("--output", type=Path, default=Path("outputs/toy_results.json"))
    args = parser.parse_args()
    result = run(args.seed, args.documents, args.queries, args.bits)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()