mixle.task.sft_plan module

sft_planner – trace-SFT for generating tool plans with parser-gated validation.

The generative rung above distill_planner(). The step-students decompose by classifying the next action; this trains a small causal LM (LM) on serialized teacher traces with the prompt-masked SFT objective (LM.fit_pairs) so the whole plan is generated:

request \n=> tool(k=v; k=v) | tool(k=v) | done \n

Free-form generation needs a strict validation boundary. The emitted text must parse under the plan grammar, every tool must exist, and every required argument must be present. Anything else escalates to the teacher and the trace is harvested. The model may emit arbitrary text, but only verified plans leave the function.

What this adds over the step-students (and what it does not): one model covers every tool with variable-length plans and can generalize over entity values it did not see during training. It does not infer unsupported tool compositions from absent traces; those cases escalate for teacher handling and future data collection.

class GenerativePlanner(lm, codec, tools, teacher, plan_agreement, max_new=160, constrained=True, conf_floor=None, lm_config=<factory>, n_requests=0, n_escalated=0, harvested=<factory>)[source]

Bases: object

A plan-writing LM behind a parse-and-validate gate: only verified plans leave; the rest escalate.

Parameters:
try_plan(request)[source]

Generate, parse, validate (grammar + specs + copy-fidelity); None = must escalate.

With constrained=True (default) the decode itself runs inside the plan grammar (mixle.task.constrained.constrained_plan_decode()): malformed text and copy-drifted values are unrepresentable, and the parse/validate below is a pure backstop.

Parameters:

request (str)

Return type:

list[dict] | None

report()[source]

Return plan agreement, escalation, and harvested-trace metrics.

Return type:

dict[str, Any]

save(path)[source]

Persist the plan-writing LM (weights + builder config), codec, specs, and gates; load() restores.

Parameters:

path (str)

Return type:

str

classmethod load(path, teacher, *, device='cpu')[source]

Reconstitute a serving GenerativePlanner from save() output plus the teacher fallback.

Parameters:
Return type:

GenerativePlanner

sft_planner(teacher, requests, tools, *, holdout=0.2, seed=0, d_model=96, n_layer=3, n_head=4, block=192, epochs=30, lr=3e-3, device='cpu', constrained=True)[source]

Trace-SFT a small causal LM into a plan writer, verified on held-out requests.

Traces serialize as request\n=> tool(k=v; ...) | ... \n pairs; LM.fit_pairs trains with the prompt masked so only plan tokens carry loss; generation stops at newline. Held-out agreement is plan-level exact match (tools + required args, in order) on requests the LM never saw.

Parameters:
Return type:

GenerativePlanner

score_plan(planner, request, plan)[source]

Mean per-character teacher-forced log-probability of a candidate plan under the trained LM.

This is not a decode: it scores a plan supplied by the CALLER (a candidate to rank against alternatives, or an already-taken plan to flag as low-probability after the fact) – the same confidence metric constrained_plan_decode() computes for its own greedy path, generalized to any plan text. Higher (less negative) is more probable; a plan scoring below the planner’s calibrated conf_floor is exactly the “low-probability plan” escalation signal used by plan-quality checks, computed explicitly here rather than left implicit in the decode loop.

Parameters:
Return type:

float

sample_plans(planner, request, n=5, *, temperature=1.0, seed=0)[source]

Draw n stochastic candidate plans from the trained LM, each scored by score_plan().

Sorted highest-score first. A draw that fails to parse or validate (the grammar is not enforced during stochastic sampling, unlike the constrained decode path) is returned as (None, -inf) – an undefined score IS the escalation signal: a generative decomposition model that cannot produce a coherent plan for a request should say so, never guess silently.

Parameters:
  • planner (GenerativePlanner)

  • request (str)

  • n (int)

  • temperature (float)

  • seed (int)

Return type:

list[tuple[list[dict] | None, float]]