Spaces:
Running on Zero
Running on Zero
| """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) | |