Fix: strip leading [BOS] token before validity-mask check and AST des…

[chaudhryumer]

Summary

Fixes two related bugs in the neural inference pipeline that caused
inference/solve.py's solve() to either produce empty output or fail
deserialization on every call, regardless of model training quality.

Bugs found and fixed

1. inference/beam_search.py — validity mask rejected every candidate token

is_valid_prefix()'s grammar parses SLaNg AST node types (NODE:TERM,
NODE:FRAC, OP:*) starting at index 0. Beam search seeds every beam
with [bos_id] before decoding starts, and passed the full token
sequence (including [BOS]) into node_pool.mask() on every step.
Since [BOS] matches none of the grammar's valid starting tokens, every
single candidate token was masked as invalid on the very first decoding
step, so safe_logits was always all -inf and generation terminated
immediately with only the seed [BOS] token.

Fix: strip the leading [BOS] token before passing to
node_pool.mask(), since it is a decoder-input framing token, not part
of the AST grammar being validated.

2. inference/solve.py — output deserialization failed on [BOS]

Same root issue at a second call site: solve() passed
result["tokens"] (still including the seed [BOS]) directly into
_verify_output(), which deserializes the token sequence as a SLaNg
AST. This failed immediately with:

Fix: strip the leading [BOS] from output_token_strings before
calling _verify_output().

Also verified while investigating

• inference/solve.py's default model_path (model/model.pkl) loads
  a stale checkpoint with an entirely different architecture (flat
  nn.Transformer, fc_out/pos_encoder keys) than the tree-based
  CalculusSolverModel actually trained by train.py
  (checkpoints/final/best.pt). Loading model/model.pkl by default
  causes load_state_dict() to fail with dozens of missing/unexpected
  key errors. Not changed in this PR (didn't want to change a public
  default without discussion), but worth a follow-up — either update
  the default model_path, or keep model/model.pkl in sync with the
  currently-trained architecture.

Testing

Verified end-to-end with a trained checkpoint from train.py:

from inference.solve import CalculusSolverInference
s = CalculusSolverInference(model_path='checkpoints/final/best.pt')
result = s.solve({
    "op": "diff", "var": "x",
    "expr": {"numi": {"terms": [{"coeff": 3, "var": {"x": 2}}]}, "deno": 1}
})

Before this fix: output_tokens: ['[BOS]'], deserialization error.
After this fix: generation proceeds past the first token and no longer
fails on [BOS] at either the validity-mask or deserialization stage.

Known follow-up

Training run so far (docs/TRAINING_RESULTS.md, 10 epochs,
max_steps=500) shows Val Rule loss plateaued at ~0.51 from epoch 1
onward with no further improvement, while Val Seq and Val Verify
dropped close to zero — likely exposure bias (teacher-forced validation
loss vs. free-running beam search generation) combined with only
~14% of the training set being seen per epoch. Currently re-running
with max_steps=3500 (full epoch coverage), learning_rate=0.0003,
and epochs=15 to see if the rule-loss plateau breaks. Will follow up
with updated docs/TRAINING_RESULTS.md / docs/EVAL_RESULTS.md once
that run completes.

[vercel]

@chaudhryumer is attempting to deploy a commit to the chuzaairforce-7557's projects Team on Vercel.

A member of the Team first needs to authorize it.