File size: 4,057 Bytes
69928d6 | 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 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | """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)
|