Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions config.json
Original file line number Diff line number Diff line change
@@ -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
Expand Down
18 changes: 16 additions & 2 deletions inference/beam_search.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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"))

Expand Down Expand Up @@ -344,4 +358,4 @@ def beam_search(
"status": status,
"root_rule_id": root_rule_id,
"root_rule_label": root_rule_label,
}
}
14 changes: 13 additions & 1 deletion inference/solve.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand Down Expand Up @@ -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()
Loading