diff --git a/inference/beam_search.py b/inference/beam_search.py index a380336..31058e2 100644 --- a/inference/beam_search.py +++ b/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 diff --git a/inference/solve.py b/inference/solve.py index 7000d08..0e218bf 100644 --- a/inference/solve.py +++ b/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__))) diff --git a/model/__init__.py b/model/__init__.py index 511ff66..997b190 100644 --- a/model/__init__.py +++ b/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", diff --git a/model/architecture.py b/model/architecture.py deleted file mode 100644 index a7e7d05..0000000 --- a/model/architecture.py +++ /dev/null @@ -1,76 +0,0 @@ -import torch -import torch.nn as nn - -from .tree_encoder import TreeEncoder -from .tree_decoder import TreeDecoder -from .rule_head import RuleHead -from .step_tracer import StepTracer - - -class CalculusModel(nn.Module): - def __init__( - self, - vocab_size, - rule_labels=None, - hidden_dim=512, - num_heads=8, - num_layers=8, - ffn_dim=2048, - dropout=0.1, - position_dim=3, - ): - super().__init__() - self.encoder = TreeEncoder( - vocab_size=vocab_size, - hidden_dim=hidden_dim, - num_heads=num_heads, - num_layers=num_layers, - ffn_dim=ffn_dim, - dropout=dropout, - position_dim=position_dim, - ) - self.rule_head = RuleHead(hidden_dim=hidden_dim, rule_labels=rule_labels) - self.decoder = TreeDecoder( - vocab_size=vocab_size, - hidden_dim=hidden_dim, - num_heads=num_heads, - num_layers=num_layers, - ffn_dim=ffn_dim, - dropout=dropout, - ) - templates = [f"Describe {label}." for label in self.rule_head.labels()] - self.step_tracer = StepTracer(hidden_dim=hidden_dim, templates=templates) - - def forward( - self, - src_tokens, - src_positions, - parent_child_pairs, - tgt_tokens, - root_mask=None, - rule_ids=None, - validity_mask=None, - src_padding_mask=None, - tgt_padding_mask=None, - memory_key_padding_mask=None, - ): - encoder_output = self.encoder( - src_tokens, src_positions, parent_child_pairs, padding_mask=src_padding_mask - ) - rule_logits = self.rule_head(encoder_output, root_mask=root_mask) - - if rule_ids is None: - rule_ids = torch.argmax(rule_logits, dim=-1) - - rule_embeddings = self.rule_head.embed_rules(rule_ids) - decoder_logits, decoder_hidden_states = self.decoder( - tgt_tokens, - encoder_output, - rule_embeddings=rule_embeddings, - validity_mask=validity_mask, - tgt_padding_mask=tgt_padding_mask, - memory_key_padding_mask=memory_key_padding_mask, - ) - - description_logits = self.step_tracer(rule_ids, decoder_hidden_states) - return decoder_logits, rule_logits, description_logits diff --git a/model/transformer.py b/model/transformer.py index df53bd4..764a0f8 100644 --- a/model/transformer.py +++ b/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 \ No newline at end of file + 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}") diff --git a/solver_model.py b/solver_model.py index beee028..cc6cc2e 100644 --- a/solver_model.py +++ b/solver_model.py @@ -1,38 +1,21 @@ +"""Compatibility shim: re-exports the one canonical model class. + +train.py and predict.py import CalculusSolverModel from here. This module used +to wrap that import in try/except and, on ImportError, silently define an +entirely different LSTM-based CalculusSolverModel -- so a broken +model/transformer.py trained the wrong network with no error raised. The +import is now unconditional: if model/transformer.py cannot be imported, this +fails loudly at import time, before any training starts. +""" + import sys -import torch -import torch.nn as nn from pathlib import Path -# Add project root directory to path safely -project_root = str(Path(__file__).parent.resolve()) -if project_root not in sys.path: - sys.path.append(project_root) +_PROJECT_ROOT = str(Path(__file__).parent.resolve()) +if _PROJECT_ROOT not in sys.path: + sys.path.append(_PROJECT_ROOT) -try: - # 🎯 FIX 2 & 3: Import official team Transformer module and check parameters signatures - from model.transformer import CalculusSolverModel - print("[Shared Architecture] Successfully hooked into the official team Transformer layout!") +# Hard import, deliberately no fallback. See module docstring. +from model.transformer import CalculusSolverModel, check_forward_contract # noqa: E402 -except (ImportError, ModuleNotFoundError): - # Fallback to structural representation with matched key routing if module is absent locally - class CalculusSolverModel(nn.Module): - def __init__(self, vocab_size=256, hidden_dim=128, num_rules=4): - super().__init__() - # Strict signature match: Team's transformer architecture handles internals directly - self.embedding = nn.Embedding(vocab_size, hidden_dim) - self.TreeEncoder = nn.LSTM(hidden_dim, hidden_dim, batch_first=True) - self.TreeDecoder = nn.LSTM(hidden_dim, hidden_dim, batch_first=True) - self.seq_generation_head = nn.Linear(hidden_dim, vocab_size) - self.RuleHead = nn.Linear(hidden_dim, num_rules) - self.StepTracer = nn.Linear(hidden_dim, 1) - - def forward(self, src_seq, tgt_in_seq): - embedded_src = self.embedding(src_seq) - enc_out, (hn, cn) = self.TreeEncoder(embedded_src) - embedded_tgt = self.embedding(tgt_in_seq) - dec_out, _ = self.TreeDecoder(embedded_tgt, (hn, cn)) - token_logits = self.seq_generation_head(dec_out) - pooled_features = enc_out[:, -1, :] - rule_logits = self.RuleHead(pooled_features) - verifier_logits = self.StepTracer(pooled_features) - return token_logits, rule_logits, verifier_logits +__all__ = ["CalculusSolverModel", "check_forward_contract"] diff --git a/tests/unit/test_model_interface.py b/tests/unit/test_model_interface.py new file mode 100644 index 0000000..173c7fd --- /dev/null +++ b/tests/unit/test_model_interface.py @@ -0,0 +1,161 @@ +"""Interface smoke tests for the model <-> training/inference boundary. + +Four separate crashes traced back to callers treating +CalculusSolverModel.forward() as returning a single tensor, and to model +classes/signatures drifting apart. These are cheap (one tiny forward pass) and +catch that class of mismatch in seconds instead of hours into a training run. +train.py runs the same check_forward_contract() before its first step. +""" + +import importlib + +import pytest + +torch = pytest.importorskip("torch") + +from model.transformer import CalculusSolverModel, check_forward_contract # noqa: E402 + +VOCAB_SIZE = 40 +NUM_RULES = 5 + + +def _tiny_model(): + torch.manual_seed(0) + return CalculusSolverModel( + vocab_size=VOCAB_SIZE, num_rules=NUM_RULES, + hidden_dim=16, num_heads=2, num_layers=1, ffn_dim=32, dropout=0.0, + ).eval() + + +def _dummy_batch(batch=3, src_len=6, tgt_len=4): + src = torch.randint(1, VOCAB_SIZE, (batch, src_len)) + tgt = torch.randint(1, VOCAB_SIZE, (batch, tgt_len)) + return src, tgt + + +# -- the 3-tuple return contract --------------------------------------------- + +def test_forward_returns_a_three_tuple(): + src, tgt = _dummy_batch() + with torch.no_grad(): + output = _tiny_model()(src, tgt) + assert isinstance(output, tuple) and len(output) == 3 + + +def test_forward_output_shapes(): + src, tgt = _dummy_batch(batch=3, tgt_len=4) + with torch.no_grad(): + decoder_logits, rule_logits, verifier_logits = _tiny_model()(src, tgt) + assert decoder_logits.shape == (3, 4, VOCAB_SIZE) + assert rule_logits.shape == (3, NUM_RULES) + assert verifier_logits.shape == (3, 1) + + +def test_forward_accepts_teacher_forced_rule_ids(): + src, tgt = _dummy_batch(batch=2) + with torch.no_grad(): + output = _tiny_model()(src, tgt, true_rule_ids=torch.tensor([0, 3])) + assert len(output) == 3 + + +def test_check_forward_contract_passes_for_the_canonical_model(): + check_forward_contract(_tiny_model(), vocab_size=VOCAB_SIZE, num_rules=NUM_RULES) + + +def test_check_forward_contract_restores_training_mode(): + model = _tiny_model().train() + check_forward_contract(model, vocab_size=VOCAB_SIZE, num_rules=NUM_RULES) + assert model.training + + +def test_check_forward_contract_rejects_a_single_tensor_return(): + class SingleTensorModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(1, 1) + + def forward(self, src_seq, tgt_in_seq): + return torch.zeros(src_seq.size(0), tgt_in_seq.size(1), VOCAB_SIZE) + + with pytest.raises(TypeError, match="3-tuple"): + check_forward_contract(SingleTensorModel(), vocab_size=VOCAB_SIZE, num_rules=NUM_RULES) + + +def test_check_forward_contract_rejects_wrong_shapes(): + model = _tiny_model() + with pytest.raises(ValueError, match="rule_logits"): + check_forward_contract(model, vocab_size=VOCAB_SIZE, num_rules=NUM_RULES + 1) + + +# -- one canonical class, no silent fallback --------------------------------- + +def test_solver_model_reexports_the_canonical_class(): + """solver_model.py used to define an LSTM model on ImportError. It must + now be the exact same class as model/transformer.py.""" + import solver_model + + assert solver_model.CalculusSolverModel is CalculusSolverModel + + +def test_solver_model_has_no_fallback_path(): + """Structural check, not a string search -- the module docstring still + describes the old fallback. No try/except and no class definitions: the + module may only re-export.""" + import ast + + source = importlib.util.find_spec("solver_model").loader.get_source("solver_model") + tree = ast.parse(source) + assert not any(isinstance(node, ast.Try) for node in ast.walk(tree)) + assert not any(isinstance(node, ast.ClassDef) for node in ast.walk(tree)) + + +def test_retired_architecture_module_is_gone(): + with pytest.raises(ModuleNotFoundError): + importlib.import_module("model.architecture") + + +# -- every model beam_search can receive honours the contract ---------------- + +def test_legacy_pkl_model_honours_the_contract(): + from inference.solve import PklTransformerModel + + model = PklTransformerModel(vocab_size=VOCAB_SIZE, hidden_dim=16).eval() + src, tgt = _dummy_batch() + with torch.no_grad(): + output = model(src, tgt, true_rule_ids=None) + assert isinstance(output, tuple) and len(output) == 3 + assert output[0].shape == (3, 4, VOCAB_SIZE) + + +def test_beam_search_runs_against_the_real_model_class(): + """End to end through beam_search with an actual CalculusSolverModel, + not a stub -- the path the tuple-index crash was found on.""" + from pathlib import Path + + from inference.beam_search import beam_search + from inference.grammar import load_vocab + + root = Path(__file__).resolve().parents[2] + vocab = load_vocab(str(root / "tokenizer" / "vocab.json")) + vocab_size = max(vocab["id_to_token"]) + 1 + num_rules = sum(1 for t in vocab["token_to_id"] if t.startswith("RULE:")) + model = CalculusSolverModel( + vocab_size=vocab_size, num_rules=num_rules, + hidden_dim=16, num_heads=2, num_layers=1, ffn_dim=32, dropout=0.0, + ).eval() + with torch.no_grad(): + result = beam_search( + model, torch.randint(1, vocab_size, (1, 8)), vocab, beam_size=2, max_len=6, + ) + assert result["status"] in ("solved", "partial") + assert isinstance(result["tokens"], list) + + +def test_beam_search_signature_has_no_compat_parameters(): + import inspect + + from inference.beam_search import beam_search + + params = inspect.signature(beam_search).parameters + assert "src_positions" not in params + assert "parent_child_pairs" not in params diff --git a/tests/unit/test_prefix_parity.py b/tests/unit/test_prefix_parity.py index bbc92ec..8d3827f 100644 --- a/tests/unit/test_prefix_parity.py +++ b/tests/unit/test_prefix_parity.py @@ -94,7 +94,14 @@ def test_training_prefix_length_matches_seed_length(self): f"This will cause positional mismatch in the decoder." ) - def test_beam_search_4_arg_model_compatibility(self): + def test_beam_search_rejects_non_canonical_model_signature(self): + """beam_search used to catch TypeError and retry with a 4-argument + form for the now-retired model/architecture.py::CalculusModel. That + also swallowed genuine TypeErrors raised inside forward(). The + 4-argument form is no longer supported: a model that does not + implement forward(src_seq, tgt_in_seq, true_rule_ids=None) must fail + loudly rather than be silently re-dispatched. + """ from inference.beam_search import beam_search, NodeValidityPool class PermissiveNodePool(NodeValidityPool): @@ -103,21 +110,16 @@ def mask(self, tokens, candidate_tokens): class MockFourArgModel(torch.nn.Module): def forward(self, src_tokens, src_positions, parent_child_pairs, tgt_tokens): - vocab_size = 120 - seq_len = tgt_tokens.size(1) - decoder_logits = torch.zeros((1, seq_len, vocab_size)) - decoder_logits[0, -1, 2] = 10.0 - rule_logits = torch.zeros((1, 13)) - return decoder_logits, rule_logits, None - - mock_model = MockFourArgModel() + raise AssertionError("must not be reached") + src_tokens = torch.tensor([[1, 10, 9, 2]]) vocab_map = { "token_to_id": self.vocab_mapping, "id_to_token": {v: k for k, v in self.vocab_mapping.items()}, } - result = beam_search(mock_model, src_tokens, vocab_map, max_len=5, node_pool=PermissiveNodePool()) - assert result["status"] == "solved" - assert len(result["tokens"]) >= 2 - + with pytest.raises(TypeError): + beam_search( + MockFourArgModel(), src_tokens, vocab_map, + max_len=5, node_pool=PermissiveNodePool(), + ) diff --git a/train.py b/train.py index dd4e807..7b7029f 100644 --- a/train.py +++ b/train.py @@ -12,7 +12,7 @@ sys.path.insert(0, os.path.abspath(os.path.dirname(__file__))) from tokenizer.slang_serializer import serialize_slang_math -from solver_model import CalculusSolverModel +from solver_model import CalculusSolverModel, check_forward_contract with open("config.json", "r") as cfg_file: config = json.load(cfg_file) @@ -358,6 +358,12 @@ def run_training_pipeline(): rule_labels=RULE_LABELS, ).to(device) + # Preflight: one dummy forward pass to confirm the (decoder_logits, + # rule_logits, verifier_logits) return contract before any real training. + # Interface mismatches used to surface only after multi-hour runs. + check_forward_contract(model, vocab_size=REAL_VOCAB_SIZE, num_rules=len(RULE_LABELS)) + print("[Preflight] Model forward contract OK", flush=True) + epochs = config.get("epochs", 1) base_lr = config["learning_rate"]