whileai.simulations.export

7 public names. import whileai.simulations as wai, then wai.name.

Name	What it does
`export_dataset`	Write `training_rows` as JSONL, gated so a broken row never reaches the trainer.
`export_preference`	Write chosen/rejected pairs as DPO-style JSONL.
`export_training`	Write `training_rows` as JSONL, gated so a broken row never reaches the trainer.
`loss_mask`	One 0/1 per message: 1 carries loss, 0 is context only.
`to_trl`	Rows from `training_rows` / `export_preference` in TRL’s shape.
`tool_call_roundtrip`	Check every tool call in exported rows carries structured arguments.
`training_rows`	Build the rows a trainer can consume directly: system prompt in, tool schemas on, wire format fixed.

export_dataset

export_dataset(
    source,
    output: str | None = None,
    system_prompt: str | None = None,
    tools: Sequence[dict] | None = None,
    strip_think: bool = True,
    validate: bool = True,
    mask_mode: str = 'assistant',
    unroll: bool = False,
    max_tool_output_chars: int | None = None,
    format: str = 'openai',
    push_to: str | None = None,
) -> dict[str, Any]

Defined in whileai/simulations/export.py. Write training_rows as JSONL, gated so a broken row never reaches the trainer. Reach for it when graded rows are ready for SFT. It builds the rows with training_rows (system prompt in, tool schemas on, tool calls in the wire format, <think> blocks stripped), checks them, and writes one JSON object per line. It returns the report: path, n and n_written, rewards, with_system, with_tools, format, mask_mode, tool_call_roundtrip, privileged_leaks, warnings, and what was cut or unrolled. export_dataset and export_training are one function object under two names (export_dataset is export_training): same arguments, same file, same report. Write export_dataset in new code (it exports a dataset, not a training run); the older spelling is kept so nothing written today breaks.

source: a SimulationData (system prompt and tools come from its profile), a row list, or a JSONL path; for the last two pass system_prompt and tools.
output: the file to write. With a path source and no output it writes <name>.train.jsonl next to it. It never overwrites the source.
validate (True): refuse to write a dataset whose tool calls do not round-trip to structured arguments, or whose assistant turns quote the row’s own privileged block (the export scrubs the key, not the reply that recited it); False exports anyway and leaves the report to read. The leak check reads the source before the scrub, so pass the SimulationData or its trajectories; rows that already came through rows(), save() or a file carry nothing to check, and report["privileged_leaks"] says so.
format: "openai" (the default) writes the OpenAI chat-completions wire row: the full messages list, function.arguments as a JSON string, and the ask carried alongside as prompt. "trl" writes what trl can actually load: conversational SFT rows (arguments as dicts, the ask under prompt_text). TRL decides “is this conversational?” from the column set, so a prompt string beside messages makes it skip the chat template without an error and train on the bare ask, which is why the TRL rows do not carry one. The report says which format and which argument encoding the round-trip gate checked.
mask_mode with format="trl": the TRL rows carry no loss_mask, because trl 0.19.1’s SFTTrainer never reads one. Its collator (DataCollatorForLanguageModeling) labels every token and unlabels only from two token-level columns the trainer builds itself: completion_mask from prompt/completion rows, and assistant_masks from assistant_only_loss=True, which needs a {% generation %} block in the chat template (Qwen2.5-Instruct has none). So mask_mode="final" and unroll=True write prompt/completion rows (prompt up to the last assistant turn, that turn as the completion), and TRL trains on exactly what the mask asked. mask_mode="assistant" (the default) writes messages rows, and TRL trains on every token of them unless assistant_only_loss=True is set. The report’s mask_mode states which of the two TRL will do, and trained_messages / masked_messages count what TRL trains, not what the SDK intended; a format="trl" export on 94 rows was measured at 9x the tokens its loss_mask marked (#507).
push_to: a Hub repo ("me/my-set") to upload the written file to, with your own token: HF_TOKEN from the environment or the login hf auth login cached, through huggingface_hub (pip install 'whileai[hf]'). Private by default; no platform call. wai.hub.push is the same upload for a file, a directory or rows you already hold, with token= and private=. The report gains hub (repo_id, url, commit).
strip_think, mask_mode, unroll, max_tool_output_chars: passed through to training_rows, which explains each.

report = wai.export_dataset(data, "train.jsonl", format="trl")
print(report["n"], report["mask_mode"], report["tool_call_roundtrip"])

export_preference

export_preference(
    pairs: Sequence[dict],
    output: str | None = None,
    system_prompt: str | None = None,
    tools: Sequence[dict] | None = None,
    strip_think: bool = True,
    validate: bool = True,
    drop_ties: bool = True,
    format: str = 'openai',
) -> dict[str, Any]

Defined in whileai/simulations/export.py. Write chosen/rejected pairs as DPO-style JSONL. With format="openai" (the default) each line is \{"prompt": "<the ask, as text>", "chosen": [...messages...], "rejected": [...messages...]\} in the same wire format as export_dataset: both sides are the whole conversation, prompt turns included, and tool-call arguments are JSON strings. With format="fireworks" each line is Fireworks’ one-turn DPO shape (input.messages is the prefix both sides share, then one assistant message each as preferred_output and non_preferred_output); pairs that lost later turns are counted as fireworks_turns_cut. With format="trl" each line is TRL’s conversational preference triple, the same one-turn preference Fireworks takes: prompt is the prefix both sides share (tool turns and later asks included) and chosen/rejected are each the one assistant turn where the sides diverge, with arguments as dicts. DPOTrainer masks the prompt and sums log-probabilities over every completion token, so a tool result or a user turn on a side would be scored as the policy’s own words; the turns after the preference are cut, and pairs that lost some are counted as trl_turns_cut. The default shape is not loadable by trl.data_utils.maybe_apply_chat_template (a prompt string with conversational sides raises TypeError: string indices must be integers), so pass format="trl" when a TRL trainer is the consumer. Pairs with no one-turn contrast (the sides never differ, or diverge on a tool result) are dropped (no_completion_dropped). The roundtrip gate runs over BOTH sides and names the encoding it checked. Pairs come from ScoredData.select_for_preference() / build_preference_pairs. A pair judge_pairs marked tie carries no preference and is left out (ties_dropped in the report) unless drop_ties=False.

export_training

export_training(
    source,
    output: str | None = None,
    system_prompt: str | None = None,
    tools: Sequence[dict] | None = None,
    strip_think: bool = True,
    validate: bool = True,
    mask_mode: str = 'assistant',
    unroll: bool = False,
    max_tool_output_chars: int | None = None,
    format: str = 'openai',
    push_to: str | None = None,
) -> dict[str, Any]

source: a SimulationData (system prompt and tools come from its profile), a row list, or a JSONL path; for the last two pass system_prompt and tools.
output: the file to write. With a path source and no output it writes <name>.train.jsonl next to it. It never overwrites the source.
validate (True): refuse to write a dataset whose tool calls do not round-trip to structured arguments, or whose assistant turns quote the row’s own privileged block (the export scrubs the key, not the reply that recited it); False exports anyway and leaves the report to read. The leak check reads the source before the scrub, so pass the SimulationData or its trajectories; rows that already came through rows(), save() or a file carry nothing to check, and report["privileged_leaks"] says so.
format: "openai" (the default) writes the OpenAI chat-completions wire row: the full messages list, function.arguments as a JSON string, and the ask carried alongside as prompt. "trl" writes what trl can actually load: conversational SFT rows (arguments as dicts, the ask under prompt_text). TRL decides “is this conversational?” from the column set, so a prompt string beside messages makes it skip the chat template without an error and train on the bare ask, which is why the TRL rows do not carry one. The report says which format and which argument encoding the round-trip gate checked.
mask_mode with format="trl": the TRL rows carry no loss_mask, because trl 0.19.1’s SFTTrainer never reads one. Its collator (DataCollatorForLanguageModeling) labels every token and unlabels only from two token-level columns the trainer builds itself: completion_mask from prompt/completion rows, and assistant_masks from assistant_only_loss=True, which needs a {% generation %} block in the chat template (Qwen2.5-Instruct has none). So mask_mode="final" and unroll=True write prompt/completion rows (prompt up to the last assistant turn, that turn as the completion), and TRL trains on exactly what the mask asked. mask_mode="assistant" (the default) writes messages rows, and TRL trains on every token of them unless assistant_only_loss=True is set. The report’s mask_mode states which of the two TRL will do, and trained_messages / masked_messages count what TRL trains, not what the SDK intended; a format="trl" export on 94 rows was measured at 9x the tokens its loss_mask marked (#507).
push_to: a Hub repo ("me/my-set") to upload the written file to, with your own token: HF_TOKEN from the environment or the login hf auth login cached, through huggingface_hub (pip install 'whileai[hf]'). Private by default; no platform call. wai.hub.push is the same upload for a file, a directory or rows you already hold, with token= and private=. The report gains hub (repo_id, url, commit).
strip_think, mask_mode, unroll, max_tool_output_chars: passed through to training_rows, which explains each.

report = wai.export_dataset(data, "train.jsonl", format="trl")
print(report["n"], report["mask_mode"], report["tool_call_roundtrip"])

loss_mask

loss_mask(messages: Sequence[dict], mode: str = 'assistant') -> list[int]

Defined in whileai/simulations/export.py. One 0/1 per message: 1 carries loss, 0 is context only. "assistant" trains every assistant turn, the multi-turn default. "final" trains only the last assistant turn, for conversations whose earlier agent turns were scripted or came from another policy (Lambert 2025, chapter Instruction Tuning, describes both). System, user, and tool messages are always 0: tool output is the environment speaking, not the policy, and training on it teaches the model to invent tool results (Lambert 2025, chapter Tool Use).

to_trl

to_trl(rows: Sequence[dict], kind: str = 'training') -> list[dict]

Defined in whileai/simulations/export.py. Rows from training_rows / export_preference in TRL’s shape. kind="training" reshapes SFT rows as conversational messages (no loss_mask: trl 0.19.1 trains on every token of them unless assistant_only_loss=True), kind="completion" as conversational prompt/completion with the last assistant turn as the completion (the mask_mode="final" mask, which the trainer honors through the completion_mask it builds), kind="preference" DPO pairs with the shared prefix as prompt and the one assistant turn where the sides diverge as each of chosen/rejected; see the module docstring for what each shape is and why it differs from the default OpenAI wire rows. Equivalent to passing format="trl" to the exporters, for callers that already hold rows. Completion rows with no assistant turn and preference pairs with no one-turn contrast (the sides never differ, or diverge on a tool result) are dropped; export_preference reports the counts.

tool_call_roundtrip

tool_call_roundtrip(
    rows: Sequence[dict],
    format: str = 'openai',
) -> dict[str, Any]

Defined in whileai/simulations/export.py. Check every tool call in exported rows carries structured arguments. Guards the training run, not the export. Which check that is depends on where the rows are going, so the report names the encoding it validated:

format="openai" (encoding: "json_string"): arguments must parse back to a dict. Arguments that survive as un-parseable strings get re-quoted by chat templates and teach the model to emit string-wrapped arguments, which then spiral on tool rejections.
format="trl" (encoding: "dict"): arguments must already be dicts. A JSON string here is valid OpenAI wire and still wrong for a chat template, which would render it quoted twice, so it counts as invalid rather than passing on a technicality.

training_rows

training_rows(
    source,
    system_prompt: str | None = None,
    tools: Sequence[dict] | None = None,
    strip_think: bool = True,
    mask_mode: str = 'assistant',
    unroll: bool = False,
    max_tool_output_chars: int | None = None,
) -> list[dict]

Defined in whileai/simulations/export.py. Build the rows a trainer can consume directly: system prompt in, tool schemas on, wire format fixed. Reach for it when you want the rows in memory rather than in a file (export_dataset writes these same rows as JSONL, with the gates). A simulated row stores the conversation without the agent’s own system prompt or tool schemas; the run knows them, the row does not. This call returns a list of dicts, one per source row, each with messages (the system message prepended, tool calls in the OpenAI wire format with id, type, function.name and function.arguments as a JSON string, each tool result linked by tool_call_id), tools (the schemas), loss_mask (which assistant turns carry loss), the ask as prompt, and the row’s grade and lineage. The module docstring has the two wire shapes.

source: a SimulationData (system prompt and tools come from its profile), a row list, or a JSONL path. For lists and paths, pass system_prompt= and tools= explicitly; a row exported without its policy trains an agent that never saw its rules.
strip_think (True): remove <think> blocks from the assistant turns, so a thinking rollout model never teaches a non-thinking student to emit them. On a reasoning base such as Qwen3 that teaches the adapter to emit an empty <think></think> and answer at once, so at eval it answers while the untrained base is still reasoning under the same max_tokens. Pass strip_think=False when the student should keep reasoning, and set thinking= the same on both arms of the eval either way.
mask_mode ("assistant"): which assistant turns carry loss (see loss_mask): all of them, or "final" for the last turn only.
unroll (False): True turns an N-turn conversation into N samples, the k-th ending at the k-th assistant turn with loss on that turn only (Lambert 2025, chapter Instruction Tuning). Every earlier agent turn then trains once with exactly the context it had, instead of only the last one (mask_mode="final") or all of them at once ("assistant", where later turns see context the policy never produced). Each sample carries unroll (turn, turns) and lineage.unrolled_from (the source row’s prompt hash and rollout index); mask_mode is ignored, and no group fields are stamped, since samples of one conversation are not a GRPO group.
max_tool_output_chars (None, cut nothing): cap each tool message at that many characters, appending [... N chars of tool output truncated] and counting the cut on the row as tool_output_truncated (messages) and tool_output_chars_cut. Tool output is masked from the loss anyway; what it costs is context, and the cut is explicit rather than silent (Lambert 2025, chapter Tool Use).

rows = wai.training_rows(data, unroll=True)
print(rows[0]["messages"][0]["role"], rows[0]["loss_mask"])

​export_dataset

​export_preference

​export_training

​loss_mask

​to_trl

​tool_call_roundtrip

​training_rows