refactor(model): one canonical model class, enforced 3-tuple forward …

inference/beam_search.py

@@ -23,33 +23,38 @@
     load_vocab,
 )
 
-def _call_model(
+def _forward(
     model,
     src_tokens: torch.Tensor,
     tgt_tokens: torch.Tensor,
-    src_positions: Optional[torch.Tensor] = None,
-    parent_child_pairs: Optional[torch.Tensor] = None,
     true_rule_ids: Optional[torch.Tensor] = None,
-) -> Any:
-    try:
-        if true_rule_ids is not None:
-            return model(src_tokens, tgt_tokens, true_rule_ids=true_rule_ids)
-        return model(src_tokens, tgt_tokens)
-    except TypeError:
-        try:
-            return model(src_tokens, tgt_tokens)
-        except TypeError:
-            device = src_tokens.device
-            batch_size, seq_len = src_tokens.size()
-            if src_positions is None:
-                src_positions = torch.zeros(
-                    (batch_size, seq_len, 3), dtype=torch.float32, device=device
-                )
-            if parent_child_pairs is None:
-                parent_child_pairs = torch.zeros(
-                    (batch_size, seq_len, seq_len), dtype=torch.float32, device=device
-                )
-            return model(src_tokens, src_positions, parent_child_pairs, tgt_tokens)
+):
+    """Call the model with the one canonical signature and unpack its output.
+
+    Every model this module accepts implements
+        forward(src_seq, tgt_in_seq, true_rule_ids=None)
+    and returns (decoder_logits, rule_logits, verifier_logits) -- see
+    model/transformer.py::CalculusSolverModel.forward.
+
+    This replaces a helper that caught TypeError and retried with other
+    argument shapes (including a 4-argument form for the retired
+    model/architecture.py::CalculusModel). That masked genuine TypeErrors
+    raised INSIDE forward() as signature mismatches. A contract violation now
+    fails immediately with a message naming the problem.
+    """
+    if true_rule_ids is None:
+        output = model(src_tokens, tgt_tokens)
+    else:
+        output = model(src_tokens, tgt_tokens, true_rule_ids=true_rule_ids)
+
+    if not isinstance(output, tuple) or len(output) != 3:
+        raise TypeError(
+            "model.forward() must return a 3-tuple "
+            "(decoder_logits, rule_logits, verifier_logits); got "
+            f"{type(output).__name__}"
+            + (f" of length {len(output)}" if isinstance(output, tuple) else "")
+        )
+    return output
 
 
 def _apply_repetition_penalty(
@@ -96,8 +101,6 @@ def beam_search(
     beam_size: int = 5,
     max_len: int = 32,
     node_pool: Optional[NodeValidityPool] = None,
-    src_positions: Optional[torch.Tensor] = None,
-    parent_child_pairs: Optional[torch.Tensor] = None,
     max_token_run: int = 4,
     no_repeat_ngram_size: int = 2,
     repetition_penalty: float = 1.2,
@@ -142,14 +145,7 @@ def beam_search(
     elif rule_token_entries:
         init_tgt = torch.tensor([[bos_id]], device=device)
         with torch.no_grad():
-            init_output = _call_model(
-                model,
-                src_tokens,
-                init_tgt,
-                src_positions=src_positions,
-                parent_child_pairs=parent_child_pairs,
-            )
-        init_rule_logits = init_output[1] if isinstance(init_output, tuple) else None
+            _, init_rule_logits, _ = _forward(model, src_tokens, init_tgt)
 
         if init_rule_logits is not None:
             pred_rule_idx = torch.argmax(init_rule_logits, dim=-1).item()
@@ -176,15 +172,9 @@ def beam_search(
 
             tgt = torch.tensor([current_tokens], device=device)
 
-            model_output = _call_model(
-                model,
-                src_tokens,
-                tgt,
-                src_positions=src_positions,
-                parent_child_pairs=parent_child_pairs,
-                true_rule_ids=true_rule_tensor,
+            decoder_logits, _, _ = _forward(
+                model, src_tokens, tgt, true_rule_ids=true_rule_tensor
             )
-            decoder_logits = model_output[0] if isinstance(model_output, tuple) else model_output
             next_logits = decoder_logits[0, -1, :]
 
             # Soft repetition penalty (task 1) -- on RAW logits, before the

inference/solve.py

@@ -7,7 +7,7 @@
 import joblib
 import torch
 
-from model.architecture import CalculusModel
+from model.transformer import CalculusSolverModel
 from inference.beam_search import NodeValidityPool, beam_search, load_vocab
 
 
@@ -30,7 +30,14 @@ def __init__(self, vocab_size: int, hidden_dim: int = 128):
         )
         self.fc_out = torch.nn.Linear(hidden_dim, vocab_size)
 
-    def forward(self, src_seq, tgt_in_seq):
+    def forward(self, src_seq, tgt_in_seq, true_rule_ids=None):
+        """Legacy model/model.pkl loader, conformed to the shared contract.
+
+        Accepts the canonical forward signature and returns the canonical
+        3-tuple, so beam_search can treat every model identically. This
+        architecture has no rule head or step tracer, so those slots are None;
+        true_rule_ids is accepted for signature parity and ignored.
+        """
         max_id = self.embedding.num_embeddings - 1
         src_seq = torch.clamp(src_seq, 0, max_id)
         tgt_in_seq = torch.clamp(tgt_in_seq, 0, max_id)
@@ -47,7 +54,7 @@ def forward(self, src_seq, tgt_in_seq):
         )
         out = self.transformer(src_emb, tgt_emb, tgt_mask=tgt_mask)
         logits = self.fc_out(out)
-        return logits
+        return logits, None, None
 
 
 class CalculusSolverInference:
@@ -83,19 +90,7 @@ def __init__(
                 vocab_size = state_dict["embedding.weight"].shape[0]
                 hidden_dim = state_dict["embedding.weight"].shape[1]
                 self.model = PklTransformerModel(vocab_size=vocab_size, hidden_dim=hidden_dim).to(self.device)
-            elif any(k.startswith(("encoder.", "decoder.", "rule_head.")) for k in state_dict.keys()) and not model_path.endswith((".pt", ".pth")):
-                self.model = CalculusModel(
-                    vocab_size=config.get("vocab_size", len(self.vocab_map["token_to_id"])),
-                    rule_labels=rule_labels,
-                    hidden_dim=config.get("hidden_dim", 512),
-                    num_heads=config.get("num_heads", 8),
-                    num_layers=config.get("num_layers", 8),
-                    ffn_dim=config.get("ffn_dim", 2048),
-                    dropout=config.get("dropout", 0.1),
-                    position_dim=config.get("position_dim", 3),
-                ).to(self.device)
             else:
-                from model.transformer import CalculusSolverModel
                 hidden_dim = 128
                 try:
                     root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

model/__init__.py

@@ -1,11 +1,12 @@
-from .architecture import CalculusModel
+from .transformer import CalculusSolverModel, check_forward_contract
 from .tree_encoder import TreeEncoder
 from .tree_decoder import TreeDecoder
 from .rule_head import RuleHead
 from .step_tracer import StepTracer
 
 __all__ = [
-    "CalculusModel",
+    "CalculusSolverModel",
+    "check_forward_contract",
     "TreeEncoder",
     "TreeDecoder",
     "RuleHead",

model/transformer.py

@@ -65,6 +65,31 @@ def __init__(
         )
 
     def forward(self, src_seq, tgt_in_seq, true_rule_ids=None):
+        """Run the full model.
+
+        RETURN CONTRACT -- always a 3-tuple, never a single tensor:
+
+            (decoder_logits, rule_logits, verifier_logits)
+
+            decoder_logits   (batch, tgt_len, vocab_size)  next-token scores
+            rule_logits      (batch, num_rules)            rule classifier
+            verifier_logits  (batch, num_templates)        step tracer; 1 template
+
+        Every caller must unpack explicitly, e.g.
+            decoder_logits, rule_logits, _ = model(src_seq, tgt_in_seq)
+        Treating the result as a tensor is what caused
+        "'tuple' object has no attribute 'reshape'" in train.py and
+        "tuple indices must be integers or slices, not tuple" in beam_search.
+
+        check_forward_contract() below enforces this; train.py runs it before
+        the first training step.
+
+        Args:
+            src_seq: (batch, src_len) source token ids.
+            tgt_in_seq: (batch, tgt_len) decoder input token ids.
+            true_rule_ids: optional (batch,) rule ids for teacher forcing the
+                rule embedding. When None the model uses argmax(rule_logits).
+        """
         device = src_seq.device
         batch_size, seq_len = src_seq.size()
 
@@ -101,4 +126,61 @@ def forward(self, src_seq, tgt_in_seq, true_rule_ids=None):
         # 5. Trace steps (verifier)
         verifier_logits = self.step_tracer(rule_ids, decoder_hidden_states)
 
-        return decoder_logits, rule_logits, verifier_logits
+        return decoder_logits, rule_logits, verifier_logits
+
+
+def check_forward_contract(
+    model: nn.Module,
+    vocab_size: int,
+    num_rules: int,
+    batch_size: int = 2,
+    src_len: int = 7,
+    tgt_len: int = 5,
+) -> None:
+    """Run one dummy forward pass and fail loudly if the contract is broken.
+
+    Checks that forward() returns exactly (decoder_logits, rule_logits,
+    verifier_logits) with the expected shapes. It costs one forward pass on a
+    tiny batch, so train.py runs it before the first step -- this class of
+    mismatch previously surfaced only after multi-hour runs.
+
+    Raises:
+        TypeError: the output is not a 3-tuple of tensors.
+        ValueError: a tensor has the wrong shape.
+    """
+    device = next(model.parameters()).device
+    was_training = model.training
+    model.eval()
+    try:
+        with torch.no_grad():
+            # Ids >= 1 so no row is all padding (pad_id 0), which would give
+            # the rule head an empty root mask.
+            src = torch.randint(1, vocab_size, (batch_size, src_len), device=device)
+            tgt = torch.randint(1, vocab_size, (batch_size, tgt_len), device=device)
+            output = model(src, tgt)
+    finally:
+        model.train(was_training)
+
+    if not isinstance(output, tuple) or len(output) != 3:
+        raise TypeError(
+            "model.forward() must return a 3-tuple "
+            "(decoder_logits, rule_logits, verifier_logits); got "
+            f"{type(output).__name__}"
+            + (f" of length {len(output)}" if isinstance(output, tuple) else "")
+        )
+
+    decoder_logits, rule_logits, verifier_logits = output
+    num_templates = getattr(getattr(model, "step_tracer", None), "num_templates", None)
+    expected = {
+        "decoder_logits": (decoder_logits, (batch_size, tgt_len, vocab_size)),
+        "rule_logits": (rule_logits, (batch_size, num_rules)),
+        "verifier_logits": (
+            verifier_logits,
+            (batch_size, num_templates) if num_templates is not None else None,
+        ),
+    }
+    for name, (tensor, shape) in expected.items():
+        if not isinstance(tensor, torch.Tensor):
+            raise TypeError(f"{name} must be a torch.Tensor, got {type(tensor).__name__}")
+        if shape is not None and tuple(tensor.shape) != shape:
+            raise ValueError(f"{name} has shape {tuple(tensor.shape)}, expected {shape}")