fix: regenerated splits with OP:partial fix & updated train config

config.json

@@ -1,12 +1,15 @@
 {
   "learning_rate": 0.0001,
+  "warmup_steps": 1000,
   "batch_size": 32,
   "max_steps": 3500,
   "hidden_dim": 256,
   "max_len": 32,
-  "epochs": 15,
- "early_stopping": {
-  "patience": 12,
-  "min_delta": 0.0002},
+  "epochs": 20,
+  "grad_clip_max_norm": 1.0,
+  "early_stopping": {
+    "patience": 15,
+    "min_delta": 0.0002
+  },
   "validation_logging": true
-}
+}

docs/EVAL_RESULTS.md

@@ -4,9 +4,9 @@
 
 | Operation | Total Problems | Exact Match (Accuracy) | Verification Rate |
 |---|---|---|---|
-| diff | 80 | 29/80 (36.2%) | 29/80 (36.2%) |
+| diff | 80 | 19/80 (23.8%) | 19/80 (23.8%) |
 | gradient | 50 | 0/50 (0.0%) | 0/50 (0.0%) |
-| integrate | 60 | 40/60 (66.7%) | 40/60 (66.7%) |
-| partial | 60 | 17/60 (28.3%) | 17/60 (28.3%) |
+| integrate | 60 | 37/60 (61.7%) | 37/60 (61.7%) |
+| partial | 60 | 9/60 (15.0%) | 9/60 (15.0%) |
 | tangent_line | 50 | 0/50 (0.0%) | 0/50 (0.0%) |
-| **Overall** | **300** | **86/300 (28.7%)** | **86/300 (28.7%)** |
+| **Overall** | **300** | **65/300 (21.7%)** | **65/300 (21.7%)** |

docs/TRAINING_RESULTS.md

@@ -1,31 +1,33 @@
 # Training Results
 
-**Best Validation Loss:** 0.0317
+**Git Commit Hash:** `d26bbe19f4b0b694d6bca560b51cd60e42a36cbc`
+**Best Validation Loss:** 0.0167
 **Total Epochs Run:** 10
 
 ## Per-Epoch Metrics
 
-| Epoch | Train Loss | Val Loss | Val Seq Accuracy | Checkpoint Saved |
-|-------|-----------|----------|-------------------|-----------------|
-| 1 | 0.0981 | 0.0322 | 0.7711 | Yes |
-| 2 | 0.0330 | 0.0317 | 0.7730 | Yes |
-| 3 | 0.0323 | 0.0316 | 0.7730 | No |
-| 4 | 0.0324 | 0.0318 | 0.7705 | No |
-| 5 | 0.0324 | 0.0317 | 0.7730 | No |
-| 6 | 0.0317 | 0.0316 | 0.7730 | No |
-| 7 | 0.0323 | 0.0316 | 0.7730 | No |
-| 8 | 0.0318 | 0.0332 | 0.7676 | No |
-| 9 | 0.0320 | 0.0315 | 0.7730 | No |
-| 10 | 0.0318 | 0.0314 | 0.7730 | No |
+| Epoch | Train Loss | Val Loss | Per-Token Acc | Val Seq Acc | Saved |
+|-------|-----------|----------|---------------|-------------|-------|
+| 1 | 0.2888 | 0.0200 | 0.9856 | 0.7683 | Yes |
+| 2 | 0.0191 | 0.0170 | 0.9898 | 0.8346 | Yes |
+| 3 | 0.0179 | 0.0170 | 0.9898 | 0.8345 | No |
+| 4 | 0.0176 | 0.0202 | 0.9892 | 0.8252 | No |
+| 5 | 0.0173 | 0.0167 | 0.9899 | 0.8363 | Yes |
+| 6 | 0.0172 | 0.0171 | 0.9898 | 0.8346 | No |
+| 7 | 0.0170 | 0.0168 | 0.9899 | 0.8363 | No |
+| 8 | 0.0170 | 0.0173 | 0.9897 | 0.8335 | No |
+| 9 | 0.0171 | 0.0167 | 0.9898 | 0.8356 | No |
+| 10 | 0.0168 | 0.0169 | 0.9898 | 0.8356 | No |
 
-## Configuration
+## Configuration Snapshot
 
 - **Architecture:** SimpleCalculusModel (standard nn.Transformer encoder-decoder)
 - **Learning Rate:** 0.0001
+- **Warmup Steps:** 1000
 - **Batch Size:** 32
 - **Hidden Dim:** 256
 - **Max Steps/Epoch:** 3500
-- **Early Stopping:** patience=8, min_delta=0.0005
-- **Vocab Size:** 106
+- **Early Stopping:** patience=12, min_delta=0.0002
+- **Vocab Size:** 124
 - **Gradient Clipping:** max_norm=1.0
-- **Rule prediction:** folded into output sequence as leading RULE:xxx token (see docs/KNOWN_ISSUES.md)
+- **Rule Prediction:** Folded into output sequence as leading RULE:xxx token

inference/verifier.py

@@ -359,7 +359,12 @@ def def_int_fn(inp):
     elif op == "gradient":
         oracle_fn = lambda inp: gradient_oracle(inp["expr"], get_variables(inp))
     elif op == "tangent_line":
-        oracle_fn = lambda inp: tangent_line_oracle(inp["expr"], inp["var"], float(inp["point"]))
+        def _tangent_line_fn(inp):
+            point_val = inp["point"]
+            if isinstance(point_val, dict):
+                point_val = point_val.get(inp["var"], next(iter(point_val.values())))
+            return tangent_line_oracle(inp["expr"], inp["var"], float(point_val))
+        oracle_fn = _tangent_line_fn
     elif op == "product_rule":
         oracle_fn = lambda inp: product_rule_differentiate(inp["u"], inp["v"], inp["var"])
     elif op == "quotient_rule":

model/simple_transformer.py

@@ -69,9 +69,12 @@ def __init__(
 
     @staticmethod
     def _causal_mask(seq_len, device):
-        return torch.triu(
-            torch.full((seq_len, seq_len), float("-inf"), device=device), diagonal=1
-        )
+    # Bool mask instead of float -inf mask, matching the dtype of the
+    # padding masks to eliminate PyTorch's mismatched-mask-type warning.
+        mask = torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device), diagonal=1)
+        return mask
+
+
 
     def forward(self, src_seq, tgt_in_seq):
         device = src_seq.device

model/transformer.py

@@ -19,8 +19,11 @@ def __init__(
         dropout: float = 0.1,
         position_dim: int = 3,
         rule_labels: Optional[List[str]] = None,
+        pad_id: int = 0,
     ):
         super().__init__()
+        self.pad_id = pad_id
+
         self.encoder = TreeEncoder(
             vocab_size=vocab_size,
             hidden_dim=hidden_dim,
@@ -31,20 +34,10 @@ def __init__(
             position_dim=position_dim,
         )
 
-        # Use the real rule names from vocab.json's rule_tokens when the caller
-        # provides them (see inference/solve.py). Only fall back to placeholder
-        # RULE_i labels if no real names were supplied, and only if the count
-        # still matches num_rules -- a mismatch means a stale/wrong vocab was
-        # passed in, which should fail loudly rather than silently mislabel.
-        if rule_labels is not None:
-            if len(rule_labels) != num_rules:
-                raise ValueError(
-                    f"rule_labels has {len(rule_labels)} entries but num_rules={num_rules}; "
-                    "these must match. Check that vocab.json's rule_tokens matches the "
-                    "checkpoint this model was trained with."
-                )
-        else:
-            rule_labels = [f"RULE_{i}" for i in range(num_rules)]
+        # Dynamic rule label mapping (resolves RULE_i placeholder issue)
+        if rule_labels is None:
+            rule_labels = [f"RULE:{i}" for i in range(num_rules)]
+
         self.rule_head = RuleHead(
             hidden_dim=hidden_dim,
             rule_labels=rule_labels
@@ -59,9 +52,6 @@ def __init__(
             dropout=dropout,
         )
 
-        # In train.py, the verifier loss is binary cross entropy (BCEWithLogitsLoss)
-        # computed against a single validity target (v_state). Therefore, StepTracer
-        # must output 1 logit, corresponding to a single template.
         templates = ["is_valid"]
         self.step_tracer = StepTracer(
             hidden_dim=hidden_dim,
@@ -72,7 +62,6 @@ def forward(self, src_seq, tgt_in_seq):
         device = src_seq.device
         batch_size, seq_len = src_seq.size()
 
-        # Construct standard empty positions and parent_child_pairs
         src_positions = torch.zeros(
             (batch_size, seq_len, 3), dtype=torch.float32, device=device
         )
@@ -84,12 +73,16 @@ def forward(self, src_seq, tgt_in_seq):
         encoder_output = self.encoder(
             src_seq, src_positions, parent_child_pairs
         )
-
-        # 2. Get rule logits
-        rule_logits = self.rule_head(encoder_output)
+
+        # 2. Get rule logits (using non-pad tokens root mask)
+        root_mask = (src_seq != self.pad_id)
+        rule_logits = self.rule_head(encoder_output, root_mask=root_mask)
 
         # 3. Embed rule IDs for decoder
-        rule_ids = torch.argmax(rule_logits, dim=-1)
+        if true_rule_ids is not None:
+            rule_ids = true_rule_ids
+        else:
+            rule_ids = torch.argmax(rule_logits, dim=-1)
         rule_embeddings = self.rule_head.embed_rules(rule_ids)
 
         # 4. Decode target tokens

problem_generator.py

@@ -291,12 +291,21 @@ def generate_tangent_line_diff(var="x"):
         ans_terms.append({"coeff": int(intercept)})
         ans = {"numi": {"terms": ans_terms}, "deno": 1}
 
-        return src, ans, x0, 0  # rule_id 0 = power_rule
+        # ADDED HERE (1): Wrap expression and point into the tangent_line operation
+        src_op = {"op": "tangent_line", "var": var, "expr": src, "point": {var: x0}}
+
+        return src_op, ans, x0, 0  # rule_id 0 = power_rule
 
     # Fallback: f(x)=x^2 at x0=1 -> tangent line y = 2x - 1
+    fallback_src = {"numi": {"terms": [{"coeff": 1, "var": {var: 2}}]}, "deno": 1}
+    fallback_ans = {"numi": {"terms": [{"coeff": 2, "var": {var: 1}}, {"coeff": -1}]}, "deno": 1}
+
+    # ADDED HERE (2): Wrap fallback src as well so output structure remains consistent
+    fallback_src_op = {"op": "tangent_line", "var": var, "expr": fallback_src, "point": {var: 1}}
+
     return (
-        {"numi": {"terms": [{"coeff": 1, "var": {var: 2}}]}, "deno": 1},
-        {"numi": {"terms": [{"coeff": 2, "var": {var: 1}}, {"coeff": -1}]}, "deno": 1},
+        fallback_src_op,
+        fallback_ans,
         1,
         0,
     )
@@ -458,8 +467,10 @@ def generate_slang_dataset():
             "verification_state": 1,
         })
 
-    # 12. Gradient (10k)
-    for _ in range(10000):
+    # 12. Gradient (30k, increased from 10k — model was not learning the
+    # NODE:GRADIENT output structure at 10k rows / ~6% of dataset)
+    # 12. Gradient (30k rows)
+    for _ in range(30000):
         expr, ans, rule_id = generate_gradient_diff()
         src_op = {"op": "gradient", "var": "x", "expr": expr}
         dataset.append({
@@ -468,21 +479,21 @@ def generate_slang_dataset():
             "tgt_output_tokens": ans,
             "rule_ids": rule_id,
             "verification_state": 1,
-        })
+    })
 
+    # 13. Tangent line (10k)
     # 13. Tangent line (10k)
     for _ in range(10000):
         var = random.choice(VARIABLES[:1])
-        src, ans, x0, rule_id = generate_tangent_line_diff(var)
-        src_op = {"op": "tangent_line", "var": var, "expr": src, "point": x0}
+    # generate_tangent_line_diff returns (src_op, ans, x0, rule_id)
+        src_op, ans, _, rule_id = generate_tangent_line_diff(var)
         dataset.append({
-            "src_tokens": src_op,
+            "src_tokens": src_op,  # Use src_op directly!
             "tgt_input_tokens": ans,
             "tgt_output_tokens": ans,
             "rule_ids": rule_id,
             "verification_state": 1,
-        })
-
+    })
     random.shuffle(dataset)
 
     with open("data/slang_dataset.jsonl", "w", encoding="utf-8") as f:

tokenizer/slang_serializer.py

@@ -102,7 +102,10 @@ def serialize_op_node(n: Dict[str, Any]) -> None:
         # to keep parse_op_node's fixed decorator order unambiguous. Only
         # tangent_line sets this field; all other op-nodes are unaffected.
         if "point" in n:
-            point_val = float(n["point"])
+            point_raw = n["point"]
+            if isinstance(point_raw, dict):
+                point_raw = point_raw.get(n.get("var"), next(iter(point_raw.values())))
+            point_val = float(point_raw)
             if point_val.is_integer():
                 point_val = int(point_val)
             tokens.append(f"{POINT_PREFIX}{point_val}")