| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from typing import Any |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| HIDDEN_DIM = 128 |
| DROPOUT = 0.0 |
|
|
|
|
| class HyperNet(nn.Module): |
| """Predict a diagonal affine transform represented by scale/bias vectors.""" |
|
|
| def __init__(self, embedding_dim: int): |
| super().__init__() |
| self.embedding_dim = embedding_dim |
| self.hypernet = nn.Sequential( |
| nn.Linear(embedding_dim, HIDDEN_DIM), |
| nn.ReLU(), |
| nn.Dropout(DROPOUT), |
| nn.Linear(HIDDEN_DIM, HIDDEN_DIM // 2), |
| nn.ReLU(), |
| nn.Dropout(DROPOUT), |
| nn.Linear(HIDDEN_DIM // 2, embedding_dim * 2), |
| ) |
| for layer in self.hypernet: |
| if isinstance(layer, nn.Linear): |
| nn.init.normal_(layer.weight, std=0.01) |
| nn.init.zeros_(layer.bias) |
|
|
| def forward(self, embedding: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| hyper_output = self.hypernet(embedding) |
| scale = torch.sigmoid(hyper_output[..., : self.embedding_dim]) |
| bias = torch.tanh(hyper_output[..., self.embedding_dim :]) |
| return scale, bias |
|
|
|
|
| @dataclass(frozen=True) |
| class ModelConfig: |
| input_dim: int |
|
|
| def __post_init__(self) -> None: |
| if isinstance(self.input_dim, bool) or not isinstance(self.input_dim, int) or self.input_dim <= 0: |
| raise ValueError("input_dim must be a positive integer.") |
|
|
| @property |
| def output_dim(self) -> int: |
| return self.input_dim |
|
|
| @property |
| def hidden_dim(self) -> int: |
| return HIDDEN_DIM |
|
|
| @property |
| def dropout(self) -> float: |
| return DROPOUT |
|
|
| def to_dict(self) -> dict[str, Any]: |
| return { |
| "input_dim": self.input_dim, |
| "output_dim": self.output_dim, |
| "hidden_dim": self.hidden_dim, |
| "dropout": self.dropout, |
| } |
|
|
|
|
| class IMRNN(nn.Module): |
| """Query/document embedding modulation for interpretable dense retrieval.""" |
|
|
| def __init__(self, config: ModelConfig): |
| super().__init__() |
| self.config = config |
| dimension = config.input_dim |
| self.projector = nn.Linear(dimension, dimension) |
| self.query_hypernet = HyperNet(dimension) |
| self.doc_hypernet = HyperNet(dimension) |
| self.query_norm = nn.LayerNorm(dimension) |
| self.doc_norm = nn.LayerNorm(dimension) |
|
|
| def project(self, embeddings: torch.Tensor) -> torch.Tensor: |
| if embeddings.shape[-1] != self.config.input_dim: |
| raise ValueError( |
| f"Expected embeddings with final dimension {self.config.input_dim}, received {embeddings.shape[-1]}." |
| ) |
| return F.normalize(self.projector(embeddings), p=2, dim=-1) |
|
|
| def _modulate( |
| self, |
| projected_queries: torch.Tensor, |
| projected_documents: torch.Tensor, |
| ) -> tuple[torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]: |
| query_scale, query_bias = self.query_hypernet(projected_queries) |
| modulated_documents = self.doc_norm(projected_documents * query_scale.unsqueeze(1) + query_bias.unsqueeze(1)) |
|
|
| document_scale, document_bias = self.doc_hypernet(projected_documents) |
| aggregate_scale = document_scale.mean(dim=1) |
| aggregate_bias = document_bias.mean(dim=1) |
| modulated_queries = self.query_norm(projected_queries * aggregate_scale + aggregate_bias) |
| details = { |
| "query_scale": query_scale, |
| "query_bias": query_bias, |
| "document_scale": document_scale, |
| "document_bias": document_bias, |
| "aggregate_document_scale": aggregate_scale, |
| "aggregate_document_bias": aggregate_bias, |
| } |
| return modulated_queries, modulated_documents, details |
|
|
| def forward_with_details( |
| self, |
| query_embeddings: torch.Tensor, |
| document_embeddings: torch.Tensor, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]: |
| if query_embeddings.dim() != 2: |
| raise ValueError("query_embeddings must have shape [batch, input_dim].") |
| if document_embeddings.dim() != 3: |
| raise ValueError("document_embeddings must have shape [batch, documents, input_dim].") |
| if query_embeddings.shape[0] != document_embeddings.shape[0]: |
| raise ValueError("Query and document batch sizes must match.") |
| if document_embeddings.shape[1] == 0: |
| raise ValueError("At least one document embedding is required.") |
|
|
| projected_queries = self.project(query_embeddings) |
| projected_documents = self.project(document_embeddings) |
| modulated_queries, modulated_documents, details = self._modulate(projected_queries, projected_documents) |
| normalized_queries = F.normalize(modulated_queries, p=2, dim=-1) |
| normalized_documents = F.normalize(modulated_documents, p=2, dim=-1) |
| scores = torch.einsum("bd,bkd->bk", normalized_queries, normalized_documents) |
| details = { |
| **details, |
| "projected_queries": projected_queries, |
| "projected_documents": projected_documents, |
| "modulated_queries": modulated_queries, |
| "modulated_documents": modulated_documents, |
| } |
| return modulated_queries, modulated_documents, scores, details |
|
|
| def forward( |
| self, |
| query_embeddings: torch.Tensor, |
| document_embeddings: torch.Tensor, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| modulated_queries, modulated_documents, scores, _ = self.forward_with_details( |
| query_embeddings, document_embeddings |
| ) |
| return modulated_queries, modulated_documents, scores |
|
|
| def score_candidates( |
| self, |
| query_embedding: torch.Tensor, |
| candidate_document_embeddings: torch.Tensor, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| if query_embedding.dim() == 1: |
| query_embedding = query_embedding.unsqueeze(0) |
| if candidate_document_embeddings.dim() == 2: |
| candidate_document_embeddings = candidate_document_embeddings.unsqueeze(0) |
| modulated_query, modulated_docs, scores = self.forward(query_embedding, candidate_document_embeddings) |
| return modulated_query.squeeze(0), modulated_docs.squeeze(0), scores.squeeze(0) |
|
|
