| """Segment-aware interactions between policy, evidence, and claim vectors.""" |
| import torch |
| from torch import Tensor, nn |
|
|
|
|
| class TriVectorInteractionHead(nn.Module): |
| """Classify policy/evidence/claim vectors using seven ordered feature blocks.""" |
|
|
| def __init__(self, hidden_dim: int, num_classes: int = 4, dropout: float = .1): |
| super().__init__() |
| for name, value in (("hidden_dim", hidden_dim), ("num_classes", num_classes)): |
| if isinstance(value, bool) or not isinstance(value, int): |
| raise TypeError(f"{name} must be an integer") |
| if value <= 0: |
| raise ValueError(f"{name} must be positive") |
| self.hidden_dim = hidden_dim |
| self.classifier = nn.Sequential( |
| nn.Dropout(dropout), |
| nn.Linear(7 * hidden_dim, 2 * hidden_dim), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Linear(2 * hidden_dim, num_classes), |
| ) |
|
|
| def construct_features(self, h_P: Tensor, h_E: Tensor, h_C: Tensor) -> Tensor: |
| """Return [P, E, C, |P-E|, P*E, |E-C|, E*C] with shape [B, 7H].""" |
| for vector in (h_P, h_E, h_C): |
| if not isinstance(vector, Tensor) or not vector.is_floating_point(): |
| raise TypeError("segment vectors must be floating-point tensors") |
| if vector.ndim != 2 or vector.shape[-1] != self.hidden_dim: |
| raise ValueError("segment vectors must have shape [B, hidden_dim]") |
| if vector.shape != h_P.shape or vector.dtype != h_P.dtype or vector.device != h_P.device: |
| raise ValueError("segment vectors must share shape, dtype, and device") |
| return torch.cat((h_P, h_E, h_C, (h_P - h_E).abs(), h_P * h_E, |
| (h_E - h_C).abs(), h_E * h_C), dim=-1) |
|
|
| def forward(self, h_P: Tensor, h_E: Tensor, h_C: Tensor) -> Tensor: |
| return self.classifier(self.construct_features(h_P, h_E, h_C)) |
|
|
| def pool_segments(self, hidden_states: Tensor, segments: Tensor, |
| attention_mask: Tensor) -> tuple[Tensor, Tensor, Tensor]: |
| """Mean-pool [B,L,H] states into policy (1), evidence (2), and claim (3). |
| |
| Segment 0 denotes delimiters/padding. Only tokens with a binary attention |
| mask of 1 contribute. Every row must contain an unmasked token in each |
| segment; empty segments raise ValueError rather than producing NaNs. |
| Float16/bfloat16 sums and division use float32 to avoid overflow and |
| count rounding; pooled vectors retain the input dtype and device. |
| """ |
| if not isinstance(hidden_states, Tensor) or not hidden_states.is_floating_point(): |
| raise TypeError("hidden_states must be a floating-point tensor") |
| if hidden_states.ndim != 3 or hidden_states.shape[-1] != self.hidden_dim: |
| raise ValueError("hidden_states must have shape [B, L, hidden_dim]") |
| for name, tensor in (("segments", segments), ("attention_mask", attention_mask)): |
| if not isinstance(tensor, Tensor): |
| raise TypeError(f"{name} must be a tensor") |
| if tensor.shape != hidden_states.shape[:2]: |
| raise ValueError(f"{name} must have shape [B, L]") |
| if tensor.device != hidden_states.device: |
| raise ValueError(f"{name} must share the hidden_states device") |
| if segments.dtype not in (torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64): |
| raise TypeError("segments must contain integer IDs") |
| if ((segments < 0) | (segments > 3)).any(): |
| raise ValueError("segment IDs must be 0, 1, 2, or 3") |
| if not ((attention_mask == 0) | (attention_mask == 1)).all(): |
| raise ValueError("attention_mask must be binary") |
| accumulation_dtype = (torch.float32 if hidden_states.dtype in (torch.float16, torch.bfloat16) |
| else hidden_states.dtype) |
| pooled = [] |
| for segment_id in (1, 2, 3): |
| selected = (segments == segment_id) & attention_mask.bool() |
| counts = selected.sum(dim=1, keepdim=True) |
| if (counts == 0).any(): |
| raise ValueError(f"empty segment {segment_id} in at least one batch row") |
| masked = hidden_states.masked_fill(~selected.unsqueeze(-1), 0) |
| mean = masked.sum(dim=1, dtype=accumulation_dtype) / counts.to(accumulation_dtype) |
| pooled.append(mean.to(hidden_states.dtype)) |
| return tuple(pooled) |
|
|
| def forward_from_hidden(self, hidden_states: Tensor, segments: Tensor, |
| attention_mask: Tensor) -> Tensor: |
| """Pool segment states and return [B, num_classes] classification logits.""" |
| return self(*self.pool_segments(hidden_states, segments, attention_mask)) |
|
|
