retraining of model

config.json

@@ -5,7 +5,7 @@
   "max_steps": 3500,
   "hidden_dim": 256,
   "max_len": 48,
-  "epochs": 5,
+  "epochs": 6,
   "grad_clip_max_norm": 1.0,
   "early_stopping": {
     "patience": 3,

docs/TRAINING_RESULTS.md

@@ -1,28 +1,19 @@
 # Training Results
 
-**Git Commit Hash:** `83e1a7266b605161bd202b58cb29a8dc55dc4be2`
-**Best Validation Loss:** 0.0045
-**Total Epochs Run:** 15
+**Git Commit Hash:** `b44635d12412ecbc8e34e1f0845a2bf3050b264f`
+**Best Validation Loss:** 0.0053
+**Total Epochs Run:** 6
 
 ## Per-Epoch Metrics
 
 | Epoch | Train Loss | Val Loss | Per-Token Acc | Val Seq Acc | Saved |
 |-------|-----------|----------|---------------|-------------|-------|
-| 1 | 0.2021 | 0.0127 | 0.9956 | 0.9418 | Yes |
-| 2 | 0.0108 | 0.0089 | 0.9966 | 0.9515 | Yes |
-| 3 | 0.0084 | 0.0069 | 0.9972 | 0.9581 | Yes |
-| 4 | 0.0073 | 0.0063 | 0.9974 | 0.9587 | Yes |
-| 5 | 0.0066 | 0.0057 | 0.9977 | 0.9623 | Yes |
-| 6 | 0.0063 | 0.0053 | 0.9977 | 0.9625 | Yes |
-| 7 | 0.0059 | 0.0053 | 0.9977 | 0.9617 | No |
-| 8 | 0.0055 | 0.0049 | 0.9977 | 0.9626 | Yes |
-| 9 | 0.0053 | 0.0047 | 0.9977 | 0.9612 | Yes |
-| 10 | 0.0056 | 0.0046 | 0.9978 | 0.9636 | No |
-| 11 | 0.0055 | 0.0046 | 0.9978 | 0.9635 | No |
-| 12 | 0.0053 | 0.0054 | 0.9974 | 0.9599 | No |
-| 13 | 0.0053 | 0.0046 | 0.9974 | 0.9563 | No |
-| 14 | 0.0053 | 0.0045 | 0.9975 | 0.9588 | Yes |
-| 15 | 0.0054 | 0.0103 | 0.9954 | 0.9435 | No |
+| 1 | 0.2052 | 0.0122 | 0.9952 | 0.9388 | Yes |
+| 2 | 0.0109 | 0.0077 | 0.9966 | 0.9517 | Yes |
+| 3 | 0.0080 | 0.0088 | 0.9970 | 0.9519 | No |
+| 4 | 0.0072 | 0.0064 | 0.9974 | 0.9597 | Yes |
+| 5 | 0.0068 | 0.0053 | 0.9977 | 0.9627 | Yes |
+| 6 | 0.0061 | 0.0051 | 0.9976 | 0.9625 | No |
 
 ## Configuration Snapshot
 

inference/beam_search.py

@@ -29,21 +29,27 @@ def _call_model(
     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:
-        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)
+        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)
 
 
 def _apply_repetition_penalty(
@@ -96,6 +102,7 @@ def beam_search(
     no_repeat_ngram_size: int = 2,
     repetition_penalty: float = 1.2,
     repetition_min_count: int = 4,
+    seed_rule_token: Optional[str] = None,
 ) -> Dict[str, Any]:
     device = src_tokens.device
     vocab = vocab_map["token_to_id"]
@@ -125,8 +132,14 @@ def beam_search(
     ]
 
     seed_tokens = [bos_id]
-
-    if rule_token_entries:
+    pred_rule_idx = None
+
+    if seed_rule_token and seed_rule_token in vocab:
+        rule_token_id = vocab[seed_rule_token]
+        seed_tokens = [bos_id, rule_token_id]
+        if seed_rule_token in rule_token_entries:
+            pred_rule_idx = rule_token_entries.index(seed_rule_token)
+    elif rule_token_entries:
         init_tgt = torch.tensor([[bos_id]], device=device)
         with torch.no_grad():
             init_output = _call_model(
@@ -146,6 +159,7 @@ def beam_search(
             if rule_token_id is not None:
                 seed_tokens = [bos_id, rule_token_id]
 
+    true_rule_tensor = torch.tensor([pred_rule_idx], device=device) if pred_rule_idx is not None else None
     beams = [{"tokens": seed_tokens, "score": 0.0, "finished": False}]
     completed = []
 
@@ -168,6 +182,7 @@ def beam_search(
                 tgt,
                 src_positions=src_positions,
                 parent_child_pairs=parent_child_pairs,
+                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, :]

inference/solve.py

@@ -163,6 +163,45 @@ def _resolve_state_dict(self, checkpoint: Any) -> Dict[str, Any]:
             return checkpoint
         raise ValueError("Unsupported checkpoint format for model state.")
 
+    def _normalize_input(self, input_env: Dict[str, Any]) -> Dict[str, Any]:
+        """Normalize bare terms or shorthand into canonical SLaNg fraction form."""
+        env = dict(input_env)
+        expr = env.get("expr")
+        if isinstance(expr, dict):
+            if "coeff" in expr and "numi" not in expr and "op" not in expr:
+                env["expr"] = {"numi": {"terms": [expr]}, "deno": 1}
+            elif "terms" in expr and "numi" not in expr:
+                env["expr"] = {"numi": expr, "deno": 1}
+        return env
+
+    def _infer_rule_token(self, input_env: Dict[str, Any]) -> Optional[str]:
+        op = input_env.get("op")
+        expr = input_env.get("expr", {})
+        if op == "integrate":
+            return "RULE:power_rule_integral"
+        if op in ("partial", "gradient"):
+            return "RULE:partial_derivative"
+        if op == "tangent_line":
+            return "RULE:power_rule"
+        if op == "diff":
+            if isinstance(expr, dict) and "op" in expr:
+                sub_op = expr["op"]
+                if sub_op in ("sin", "cos", "tan"):
+                    return "RULE:trig_rule"
+                if sub_op == "exp":
+                    return "RULE:exp_rule"
+                if sub_op == "ln":
+                    return "RULE:log_rule"
+            terms = expr.get("numi", {}).get("terms", []) if isinstance(expr, dict) else []
+            if len(terms) > 1:
+                return "RULE:sum_rule"
+            if len(terms) == 1:
+                t = terms[0]
+                if not t.get("var") or all(v == 0 for v in t.get("var", {}).values()):
+                    return "RULE:constant_rule"
+                return "RULE:power_rule"
+        return None
+
     def _serialize_input(self, input_env: Dict[str, Any]) -> List[str]:
         from tokenizer.slang_serializer import serialize_slang_math
         return serialize_slang_math(input_env)
@@ -174,7 +213,9 @@ def _verify_output(
         return verify(input_env, output_tokens)
 
     def solve(self, input_env: Dict[str, Any]) -> Dict[str, Any]:
-        token_strings = self._serialize_input(input_env)
+        normalized_env = self._normalize_input(input_env)
+        rule_token = self._infer_rule_token(normalized_env)
+        token_strings = self._serialize_input(normalized_env)
         token_ids = [
             self.vocab_map["token_to_id"].get(token, self.pad_id)
             for token in token_strings
@@ -191,6 +232,7 @@ def solve(self, input_env: Dict[str, Any]) -> Dict[str, Any]:
             beam_size=self.beam_size,
             max_len=self.max_len,
             node_pool=self.node_pool,
+            seed_rule_token=rule_token,
         )
 
         output_token_strings = [
@@ -207,7 +249,7 @@ def solve(self, input_env: Dict[str, Any]) -> Dict[str, Any]:
             predicted_rule = output_token_strings[0]
             output_token_strings = output_token_strings[1:]
 
-        verifier_result = self._verify_output(input_env, output_token_strings)
+        verifier_result = self._verify_output(normalized_env, output_token_strings)
         if verifier_result.get("status") in ("solved", "unverified", "unsolvable"):
             result["status"] = verifier_result["status"]
         result["verified"] = verifier_result.get("verified", False)

interactive_cli.py

@@ -2,6 +2,12 @@
 import json
 from pathlib import Path
 
+if sys.stdout.encoding != 'utf-8':
+    try:
+        sys.stdout.reconfigure(encoding='utf-8')
+    except Exception:
+        pass
+
 ROOT = Path(__file__).resolve().parent
 if str(ROOT) not in sys.path:
     sys.path.insert(0, str(ROOT))
@@ -14,8 +20,8 @@ def run_cli():
     print("=" * 60)
 
     candidate_paths = [
-        ROOT / "model" / "model.pkl",
         ROOT / "checkpoints" / "final" / "best.pt",
+        ROOT / "model" / "model.pkl",
         ROOT / "checkpoints" / "checkpoint_epoch_1.pt",
     ]
     checkpoint_path = None
@@ -40,8 +46,8 @@ def run_cli():
     print(" 1) d/dx (3x^2)")
     print(" 2) d/dx (5x^3 + 2x)")
     print(" 3) d/dx (sin(x))")
-    print(" 4) ∫ (x^2) dx")
-    print(" 5) ∂/∂x (x^2 * y^3)\n")
+    print(" 4) integrate (x^2) dx")
+    print(" 5) partial d/dx (x^2 * y^3)\n")
 
     while True:
         try:
@@ -56,15 +62,15 @@ def run_cli():
 
         expr_dict = None
         if user_input == "1":
-            expr_dict = {"op": "diff", "var": "x", "expr": {"coeff": 3, "var": {"x": 2}}}
+            expr_dict = {"op": "diff", "var": "x", "expr": {"numi": {"terms": [{"coeff": 3, "var": {"x": 2}}]}, "deno": 1}}
         elif user_input == "2":
             expr_dict = {"op": "diff", "var": "x", "expr": {"numi": {"terms": [{"coeff": 5, "var": {"x": 3}}, {"coeff": 2, "var": {"x": 1}}]}, "deno": 1}}
         elif user_input == "3":
-            expr_dict = {"op": "diff", "var": "x", "expr": {"op": "sin", "arg": {"var": {"x": 1}}}}
+            expr_dict = {"op": "diff", "var": "x", "expr": {"op": "sin", "expr": {"numi": {"terms": [{"coeff": 1, "var": {"x": 1}}]}, "deno": 1}}}
         elif user_input == "4":
-            expr_dict = {"op": "integrate", "var": "x", "expr": {"var": {"x": 2}}}
+            expr_dict = {"op": "integrate", "var": "x", "expr": {"numi": {"terms": [{"coeff": 3, "var": {"x": 2}}]}, "deno": 1}}
         elif user_input == "5":
-            expr_dict = {"op": "partial", "var": "x", "expr": {"coeff": 1, "var": {"x": 2, "y": 3}}}
+            expr_dict = {"op": "partial", "var": "x", "expr": {"numi": {"terms": [{"coeff": 1, "var": {"x": 2, "y": 3}}]}, "deno": 1}}
         else:
             try:
                 expr_dict = json.loads(user_input)

predict.py

@@ -3,6 +3,12 @@
 import torch
 from pathlib import Path
 
+if sys.stdout.encoding != 'utf-8':
+    try:
+        sys.stdout.reconfigure(encoding='utf-8')
+    except Exception:
+        pass
+
 sys.path.insert(0, str(Path(__file__).parent.resolve()))
 
 from solver_model import CalculusSolverModel
@@ -89,8 +95,10 @@ def evaluate_cli_input():
         rule_labels=RULE_LABELS,
     )
 
-    checkpoint_path = "checkpoints/checkpoint_epoch_1.pt"
-    state_dict = torch.load(checkpoint_path, map_location="cpu")
+    checkpoint_path = Path("checkpoints/final/best.pt")
+    if not checkpoint_path.exists():
+        checkpoint_path = Path("checkpoints/checkpoint_epoch_1.pt")
+    state_dict = torch.load(str(checkpoint_path), map_location="cpu")
     model.load_state_dict(state_dict)  # let mismatches raise loudly, never swallow them
     model.eval()
 

train.py

@@ -19,7 +19,6 @@
 
 
 def get_git_commit_hash():
-    """Returns the exact current git commit hash for provenance tracking."""
     try:
         hash_str = subprocess.check_output(["git", "rev-parse", "HEAD"]).decode("utf-8").strip()
         return hash_str
@@ -28,10 +27,6 @@ def get_git_commit_hash():
 
 
 def flatten_vocab(raw_vocab):
-    """
-    Same flattening rule as inference/beam_search.flatten_vocab on org main:
-    merge every sub-dict, skip keys starting with '_' (e.g. _comment, _version).
-    """
     flat = {}
     for key, value in raw_vocab.items():
         if key.startswith("_"):
@@ -46,10 +41,8 @@ def flatten_vocab(raw_vocab):
 
 vocab_mapping = flatten_vocab(_raw_vocab)
 
-# IDs are NOT contiguous (gaps by design — see docs/KNOWN_ISSUES.md, STRUCT:OPEN @ 23).
 REAL_VOCAB_SIZE = max(vocab_mapping.values()) + 1
 
-# Rule labels/tokens, derived from vocab's rule_tokens, ordered by ID.
 _rule_items = sorted(_raw_vocab.get("rule_tokens", {}).items(), key=lambda kv: kv[1])
 RULE_LABELS = [name.split("RULE:", 1)[1] for name, _ in _rule_items]
 RULE_TOKEN_STRINGS = [name for name, _ in _rule_items]
@@ -140,7 +133,7 @@ def preflight_check_max_len(dataset_path, max_len):
         print(f"[Pre-flight] OK: max_len has {max_len - worst} tokens of headroom.")
 
 
-def evaluate_validation(model, val_loader, criterion):
+def evaluate_validation(model, val_loader, criterion, device="cpu"):
     model.eval()
     total_loss = 0.0
     total_correct_seq = 0
@@ -151,10 +144,10 @@ def evaluate_validation(model, val_loader, criterion):
 
     with torch.no_grad():
         for batch in val_loader:
-            src_seq = batch["src_seq"]
-            tgt_in = batch["tgt_in_seq"][:, :-1]
-            tgt_out = batch["tgt_out_seq"][:, 1:]
-            rule_id = batch["rule_id"]
+            src_seq = batch["src_seq"].to(device)
+            tgt_in = batch["tgt_in_seq"][:, :-1].to(device)
+            tgt_out = batch["tgt_out_seq"][:, 1:].to(device)
+            rule_id = batch["rule_id"].to(device)
 
             decoder_logits, rule_logits, verifier_logits = model(src_seq, tgt_in, true_rule_ids=rule_id)
             loss = criterion(decoder_logits.reshape(-1, REAL_VOCAB_SIZE), tgt_out.reshape(-1))
@@ -354,13 +347,17 @@ def run_training_pipeline():
         print(f"[DEBUG] Val dataset loaded: {len(val_dataset)} examples", flush=True)
         val_loader = DataLoader(val_dataset, batch_size=config["batch_size"], shuffle=False)
 
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    device_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU"
+    print(f"[Hardware] Training device: {device} ({device_name})", flush=True)
+
     print("[DEBUG] Building model...", flush=True)
     model = CalculusSolverModel(
         vocab_size=REAL_VOCAB_SIZE,
         num_rules=len(RULE_LABELS),
         hidden_dim=config["hidden_dim"],
         rule_labels=RULE_LABELS,
-    )
+    ).to(device)
 
     # Move 'epochs' extraction UP before scheduler calculations
     epochs = config.get("epochs", 1)
@@ -428,10 +425,10 @@ def lr_lambda(current_step):
                 print(f"[DEBUG] epoch {epoch} step {step} - batch received, running forward/backward...", flush=True)
             optimizer.zero_grad()
 
-            src_seq = batch["src_seq"]
-            tgt_in_full = batch["tgt_in_seq"]
-            tgt_out = batch["tgt_out_seq"][:, 1:]
-            rule_id = batch["rule_id"]
+            src_seq = batch["src_seq"].to(device)
+            tgt_in_full = batch["tgt_in_seq"].to(device)
+            tgt_out = batch["tgt_out_seq"][:, 1:].to(device)
+            rule_id = batch["rule_id"].to(device)
 
             tgt_in = tgt_in_full[:, :-1]
 
@@ -448,7 +445,6 @@ def lr_lambda(current_step):
                 ss_mask = ss_mask & ~special_mask
 
                 tgt_in = torch.where(ss_mask, pred_tokens, tgt_in)
-
             decoder_logits, rule_logits, verifier_logits = model(src_seq, tgt_in, true_rule_ids=rule_id)
             loss = criterion(decoder_logits.reshape(-1, REAL_VOCAB_SIZE), tgt_out.reshape(-1))
 
@@ -476,7 +472,7 @@ def lr_lambda(current_step):
         }
 
         if val_loader is not None:
-            val_loss, val_seq_acc, val_token_acc = evaluate_validation(model, val_loader, criterion)
+            val_loss, val_seq_acc, val_token_acc = evaluate_validation(model, val_loader, criterion, device=device)
 
             num_proxy_examples = config.get("proxy_eval_examples", 15)
             fr_seq_acc, fr_token_acc, fr_avg_len = evaluate_free_running(