| """Attention reasoning on soft symbols; perception and task heads are supplied.""" |
|
|
| from typing import NamedTuple |
|
|
| import torch |
| from torch import nn |
|
|
| from softasp_solver.components.evidence import EvidenceMode, clamp_evidence, restore_evidence |
|
|
|
|
| class ReasoningOutput(NamedTuple): |
| pre_logits: torch.Tensor |
| logits: torch.Tensor |
|
|
|
|
| def build_attention( |
| d_model: int, n_heads: int, n_layers: int, dropout: float |
| ) -> nn.TransformerEncoder: |
| """Construct the existing pre-norm/GELU stack without extra initialization.""" |
| layer = nn.TransformerEncoderLayer( |
| d_model=d_model, |
| nhead=n_heads, |
| dim_feedforward=d_model * 4, |
| dropout=dropout, |
| activation="gelu", |
| batch_first=True, |
| norm_first=True, |
| ) |
| return nn.TransformerEncoder(layer, num_layers=n_layers) |
|
|
|
|
| def reason_from_probabilities( |
| probabilities: torch.Tensor, |
| *, |
| projection: nn.Module, |
| structure: torch.Tensor | None, |
| attention: nn.Module, |
| head: nn.Module, |
| ) -> torch.Tensor: |
| """Compose probability projection, optional structure, attention and task head. |
| |
| Structure can be (N,D) or (B,N,D). A head can produce per-position or pooled |
| task outputs; the framework imposes no Sudoku vocabulary or board shape. |
| """ |
| tokens = projection(probabilities) |
| if structure is not None: |
| tokens = tokens + (structure.unsqueeze(0) if structure.ndim == 2 else structure) |
| return head(attention(tokens)) |
|
|
|
|
| class AS2Network(nn.Module): |
| """Reusable composition for an instantiation supplying its own encoder/head. |
| |
| The encoder returns (B,N,D) embeddings. Objectives are separate modules. |
| The Sudoku compatibility adapter uses the same functions with its historical |
| parameter ownership, avoiding checkpoint renaming or duplicate registration. |
| """ |
|
|
| def __init__( |
| self, |
| *, |
| encoder: nn.Module, |
| concept_head: nn.Module, |
| projection: nn.Module, |
| attention: nn.Module, |
| head: nn.Module, |
| structure: nn.Module | None = None, |
| training_evidence: EvidenceMode = "observed", |
| inference_evidence: EvidenceMode = "predicted", |
| restore_observed: bool = False, |
| ) -> None: |
| super().__init__() |
| if inference_evidence == "observed": |
| raise ValueError("Inference cannot require target symbols") |
| self.encoder = encoder |
| self.concept_head = concept_head |
| self.projection = projection |
| self.attention = attention |
| self.head = head |
| self.structure = structure |
| self.training_evidence: EvidenceMode = training_evidence |
| self.inference_evidence: EvidenceMode = inference_evidence |
| self.restore_observed = restore_observed |
|
|
| def forward( |
| self, |
| inputs: torch.Tensor, |
| evidence_mask: torch.Tensor | None = None, |
| observed_symbols: torch.Tensor | None = None, |
| structure: torch.Tensor | None = None, |
| ) -> ReasoningOutput: |
| if not self.training and observed_symbols is not None: |
| raise ValueError("Inference must not receive target symbols") |
| pre_logits = self.concept_head(self.encoder(inputs)) |
| probabilities = clamp_evidence( |
| pre_logits, |
| evidence_mask, |
| mode=self.training_evidence if self.training else self.inference_evidence, |
| observed_symbols=observed_symbols, |
| ) |
| if structure is None and self.structure is not None: |
| structure = self.structure() |
| logits = reason_from_probabilities( |
| probabilities, |
| projection=self.projection, |
| structure=structure, |
| attention=self.attention, |
| head=self.head, |
| ) |
| if not self.training and self.restore_observed: |
| if evidence_mask is None: |
| raise ValueError("Evidence restoration requires an evidence mask") |
| logits = restore_evidence(logits, pre_logits, evidence_mask) |
| return ReasoningOutput(pre_logits, logits) |
|
|
