Custom Dataset Creation

This guide shows how to prepare custom datasets for ATLAS training. For dataset references and schemas, see Datasets Reference.

Data Format Requirements

Your dataset should follow this structure:

{
    "prompt": str,           # Required: Task description
    "ground_truth": str,     # Required: Correct solution
    "metadata": {            # Optional: Additional context
        "difficulty": str,
        "source": str,
        "tags": List[str]
    }
}

Preprocessing Pipeline (JSONL exports)

Use the runtime helpers that ship in this repository to turn SDK exports into trainer-ready splits:

from transformers import AutoTokenizer
from atlas_core.data.runtime_traces import get_runtime_trace_dataset

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct")
splits = get_runtime_trace_dataset(
    tokenizer=tokenizer,
    export_path="traces/runtime.jsonl",  # Generated via `arc-atlas export …`
    eval_split_ratio=0.1,
    dataset_max_samples=5000,
)

train_ds = splits["train_dataset"]
eval_ds = splits["eval_dataset"]

Expected output:

Loading dataset from traces/runtime.jsonl...
Loaded 5000 examples
Creating train/eval splits (90/10)...
Train dataset: 4500 examples
Eval dataset: 500 examples

Postgres-Backed Workflows

For Postgres-backed workflows, query the SDK database directly and convert records with atlas_core.data.runtime_traces:

from atlas.training_data import get_training_sessions
from atlas_core.data.runtime_traces import sessions_to_rl_records

sessions = get_training_sessions(
    db_url="postgresql://atlas:atlas@localhost:5433/atlas",
    min_reward=0.8,
    learning_key="security-review",
)

records = sessions_to_rl_records(sessions)

records is a list of dictionaries that any Hugging Face Dataset can ingest (the same structure Hydra configs consume via atlas_core.data.runtime_traces). See the Training Data Pipeline guide for additional filters and batching helpers.

GKD alignment note: Every conversation record now carries prompt_text (serialized messages excluding the final assistant turn) and completion_text (the assistant response the student learns to mimic). These fields let the distillation pipeline re-render prompts with both the student and teacher tokenizers so cross-tokenizer KL is computed in each model’s native chat template.

Quality Validation

Inspect coverage with standard Python tooling—you already have datasets installed for training:

from collections import Counter
from datasets import Dataset

dataset = Dataset.from_list(records)
lengths = [len(example["step_trace"].split()) for example in dataset]
domains = Counter(example["session_metadata"].get("domain", "unknown") for example in dataset)

print(f"Examples: {len(dataset)}")
print(f"Avg step length: {sum(lengths)/len(lengths):.1f} tokens")
print(f"Domains: {domains}")

Expected output:

Examples: 5000
Avg step length: 324.5 tokens
Domains: Counter({'math': 2100, 'code': 1800, 'debug': 700, 'reasoning': 400})

Pair these quick checks with any in-house validators your team already maintains. The key is to keep the format (prompt, student_response, guidance, rewards) identical to what the SDK emits so Atlas Core can reuse the traces without custom glue code.

Next Steps

Datasets Reference

Official datasets and schemas

GKD Training

Train with custom datasets

GRPO Training

RL training with custom data

Export Traces

Export runtime traces from the SDK

Getting Started

SDK Guides

Examples

Training

Core Concepts

Reference

Benchmarks

Custom Dataset Creation

Data Format Requirements

Preprocessing Pipeline (JSONL exports)

Postgres-Backed Workflows

Quality Validation

Next Steps

Datasets Reference

GKD Training

GRPO Training

Export Traces

​Data Format Requirements

​Preprocessing Pipeline (JSONL exports)

​Postgres-Backed Workflows

​Quality Validation

​Next Steps

Datasets Reference

GKD Training

GRPO Training

Export Traces

Data Format Requirements

Preprocessing Pipeline (JSONL exports)

Postgres-Backed Workflows

Quality Validation

Next Steps