dataset

environment Mar 19, 2026 3 min read

the dataset is what your model trains on. during training, each example becomes one rollout: the model receives the example’s messages (a single user message, or a full conversation history), generates a response, and gets scored by your reward function.

format

training data is a list of dicts. the simplest format uses prompt and ground_truth columns:

train_data = [
    {"prompt": "what is 2+2?", "ground_truth": "4"},
    {"prompt": "capital of France?", "ground_truth": "Paris"},
]

for multi-turn conversations, use messages instead of prompt:

train_data = [
    {
        "messages": [
            {"role": "user", "content": "search for climate change"},
            {"role": "assistant", "content": "searching..."},
            {"role": "user", "content": "summarize the top result"},
        ],
        "ground_truth": "...",
    },
]

each message is a Message dict. the web UI’s dataset viewer reads the same structure from a JSONL file, one {"messages": [...]} object per line.

Message

every message, whether in a messages column or in the resulting Example (see preprocessing), is a Message from benchmax.envs (an OpenAI chat-completion message dict):

fieldtypedescription
rolestr"system", "user", "assistant", or "tool" (required)
contentstrthe message text
tool_callslist[ToolCallDict]tool invocations (assistant messages). serialized in OpenAI nested format
tool_call_idstrwhich tool call this message responds to (tool messages)
namestroptional participant name (system/user/assistant messages)

where training data comes from

  • manual creation: write examples by hand. good for small, focused tasks.
  • synthetic generation: use the QA generation pipeline to create examples from a corpus. good for RAG tasks.
  • public datasets: resolve an existing dataset inside create_dataset: the geo3k example loads a hugging face dataset at runtime.

custom preprocessing

reshape rows in your environment’s row_to_example callback when your column names differ from the defaults:

from pathlib import Path

from benchmax.envs import BaseEnv, Example, JsonlDataset, canonical_example_id

class MyEnv(BaseEnv):
    async def create_dataset(self, split, base_dir: Path, *, max_examples=None):
        return JsonlDataset(
            base_dir / f"{split}.jsonl",
            row_to_example=self._example_from_row,
            max_examples=max_examples,
        )

    def _example_from_row(self, row: dict) -> Example:
        payload = {
            "prompt_messages": [{"role": "user", "content": row["question"]}],
            "ground_truth": row.get("answer"),
        }
        return Example(id=canonical_example_id(payload), payload=payload)

canonical_example_id hashes the payload’s JSON content (independent of key order), so the same data always gets the same id. Example has these fields:

fieldtypedescription
idstrstable identity for the example, computed with canonical_example_id(payload)
payloaddict[str, Any]the example data. for BaseEnv, must contain prompt_messages (a list[Message]); every other key is opaque environment data for your hooks

per-example setup data goes in the payload too. for instance, a spreadsheet environment might pass the path to the input file:

payload = {
    "prompt_messages": [{"role": "user", "content": "fix the formula in cell B5"}],
    "answer_position": "B5",
    "expected": "42",
    "spreadsheet_path": "/data/input.xlsx",
}
return Example(id=canonical_example_id(payload), payload=payload)

guidance

  • minimum size: 200+ examples are recommended for meaningful training, 1000+ for best results.
  • validate before launching. run a few examples through your row_to_example locally to catch format issues. validate_environment does this for you - see testing.