From 83bfc7ec74ce1b76f3f67938a175b59aef526f21 Mon Sep 17 00:00:00 2001
From: chaudhryumer <chumerkhalil105@gmail.com>
Date: Wed, 22 Jul 2026 01:04:37 +0500
Subject: [PATCH] Fix: strip leading [BOS] token before validity-mask check and
 AST deserialization

---
 config.json              |  7 ++++---
 inference/beam_search.py | 18 ++++++++++++++++--
 inference/solve.py       | 14 +++++++++++++-
 3 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/config.json b/config.json
index 6eec069..3e2e380 100644
--- a/config.json
+++ b/config.json
@@ -1,9 +1,10 @@
 {
-  "learning_rate": 0.0001,
+  "learning_rate": 0.0003,
   "batch_size": 32,
-  "max_steps": 500,
+  "max_steps": 3500,
   "hidden_dim": 128,
-  "epochs": 10,
+  "max_len": 32,
+  "epochs": 15,
   "early_stopping": {
     "patience": 5,
     "min_delta": 0.001
diff --git a/inference/beam_search.py b/inference/beam_search.py
index 76200d8..02aedc4 100644
--- a/inference/beam_search.py
+++ b/inference/beam_search.py
@@ -281,6 +281,20 @@ def beam_search(
 
             current_tokens = beam["tokens"]
             token_strings = [id_to_token[token_id] for token_id in current_tokens]
+
+            # FIX (docs/KNOWN_ISSUES.md): is_valid_prefix()'s grammar parses SLaNg
+            # AST structure only — it has no rule for a leading [BOS] token, so
+            # passing token_strings as-is caused every candidate to be marked
+            # invalid on the very first decoding step, producing empty output
+            # ([BOS] only) on every solve() call regardless of training quality.
+            # Strip the seed [BOS] before validity checking; it is not part of
+            # the AST grammar being validated.
+            validity_tokens = (
+                token_strings[1:]
+                if token_strings and token_strings[0] == "[BOS]"
+                else token_strings
+            )
+
             tgt = torch.tensor([current_tokens], device=device)
             decoder_logits, _ = model.decoder(
                 tgt,
@@ -291,7 +305,7 @@ def beam_search(
                 memory_key_padding_mask=None,
             )
             next_logits = decoder_logits[0, -1, :]
-            mask = node_pool.mask(token_strings, all_candidate_tokens)
+            mask = node_pool.mask(validity_tokens, all_candidate_tokens)
             invalid_mask = torch.tensor([not valid for valid in mask], device=device)
             safe_logits = next_logits.masked_fill(invalid_mask, float("-inf"))
 
@@ -344,4 +358,4 @@ def beam_search(
         "status": status,
         "root_rule_id": root_rule_id,
         "root_rule_label": root_rule_label,
-    }
+    }
\ No newline at end of file
diff --git a/inference/solve.py b/inference/solve.py
index efbcbb8..0aa8869 100644
--- a/inference/solve.py
+++ b/inference/solve.py
@@ -161,6 +161,18 @@ def solve(self, input_env: Dict[str, Any]) -> Dict[str, Any]:
             if token_id in self.vocab_map["id_to_token"]
         ]
 
+        # FIX (docs/KNOWN_ISSUES.md): beam_search seeds every beam with a
+        # leading [BOS] token, which is correct for decoder input framing but
+        # is not part of the SLaNg AST grammar itself. Downstream consumers
+        # (the deserializer inside verify(), and any AST-structure parsing)
+        # expect a pure token sequence starting at a real node type
+        # (NODE:TERM / NODE:FRAC / OP:...), not [BOS]. Without this strip,
+        # deserialization fails immediately with "Unexpected token ... [BOS]"
+        # on every single call, regardless of whether the underlying sequence
+        # the model generated was otherwise valid.
+        if output_token_strings and output_token_strings[0] == "[BOS]":
+            output_token_strings = output_token_strings[1:]
+
         verifier_result = self._verify_output(input_env, output_token_strings)
         if verifier_result.get("status") in ("solved", "unverified", "unsolvable"):
             result["status"] = verifier_result["status"]
@@ -195,4 +207,4 @@ def solve(self, input_env: Dict[str, Any]) -> Dict[str, Any]:
     try:
         print(json.dumps(solver.solve(payload), indent=2))
     finally:
-        solver.close()
+        solver.close()
\ No newline at end of file