| """MLP fusion of frozen signal, question and candidate embeddings.""" | |
| import torch | |
| from torch import nn | |
| class MLPFusionHead(nn.Module): | |
| """Score candidates with an MLP and Hadamard product features.""" | |
| def __init__(self, width=128, dropout=0.1): | |
| super().__init__() | |
| if width < 2 or not 0 <= dropout < 1: | |
| raise ValueError("Invalid MLP dimensions or dropout") | |
| self.config = dict(kind="interaction", width=width, dropout=dropout) | |
| self.signal = nn.Sequential(nn.LayerNorm(512), nn.Linear(512, width), nn.GELU()) | |
| self.text = nn.Sequential(nn.Linear(768, width), nn.LayerNorm(width)) | |
| self.context = nn.Sequential(nn.Linear(width * 3, width), nn.GELU(), nn.Dropout(dropout)) | |
| self.score = nn.Sequential(nn.Linear(width * 3, width), nn.GELU(), | |
| nn.Dropout(dropout), nn.Linear(width, 1)) | |
| def forward(self, signal, signal_mask, question, criteria, criteria_mask): | |
| batch = len(signal) | |
| if signal.shape != (batch, 1, 512) or signal_mask.shape != (batch, 1) or not signal_mask.all(): | |
| raise ValueError("MLP head requires one real signal vector") | |
| if question.shape != (batch, 768) or criteria.ndim != 3 or criteria.shape[0] != batch or criteria.shape[-1] != 768: | |
| raise ValueError("Question and criteria must have separate 768-dimensional vectors") | |
| if (criteria_mask.shape != criteria.shape[:2] or criteria_mask.dtype != torch.bool | |
| or signal_mask.dtype != torch.bool or not (criteria_mask.sum(-1) >= 2).all()): | |
| raise ValueError("Need boolean masks and at least two real criteria") | |
| if not all(torch.isfinite(x).all() for x in (signal, question, criteria)): | |
| raise ValueError("Head inputs must be finite") | |
| x, q, c = self.signal(signal[:, 0]), self.text(question), self.text(criteria) | |
| context = self.context(torch.cat([x, q, x * q], -1))[:, None].expand_as(c) | |
| z = self.score(torch.cat([context, c, context * c], -1)).squeeze(-1) | |
| if not torch.isfinite(z[criteria_mask]).all(): | |
| raise FloatingPointError("Nonfinite decision logits") | |
| return z.masked_fill(~criteria_mask, -torch.inf) | |