File size: 4,802 Bytes
a82e638 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | """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))
|