src/visual_doc_rag/maxsim.py

"""ColBERT late interaction: sum each query token's best document-token match."""

import torch


def maxsim(query: torch.Tensor, documents: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    """Score Q×D query against B×P×D pages; mask marks real page tokens.

    Mask padding with -inf, rather than zero: valid similarities can be negative.
    Inputs are model-normalized embeddings; accumulate scores in float32.
    """
    if query.ndim != 2 or documents.ndim != 3 or query.shape[-1] != documents.shape[-1]:
        raise ValueError("Expected query [Q,D] and documents [B,P,D] with matching D.")
    if mask.shape != documents.shape[:2] or mask.dtype != torch.bool:
        raise ValueError("Expected boolean document mask [B,P].")
    if query.shape[0] == 0 or not mask.any(dim=1).all():
        raise ValueError("Query and each document must contain at least one token.")
    similarities = torch.einsum("qd,bpd->bqp", query.float(), documents.float())
    return similarities.masked_fill(~mask[:, None, :], -torch.inf).amax(dim=-1).sum(dim=-1)