DataOrchestra: Learning to Orchestrate Per-Example Curation of Pretraining Data

📄 Paper   |   🤗 Orchestrator Model   |   🤗 NP Model  

This repository holds the open-source code for the paper DataOrchestra: Learning to Orchestrate Per-Example Curation of Pretraining Data. It releases the full framework: the pipeline that builds the orchestrator's training data, the scripts that run the orchestrator together with its downstream tool models, and all the prompts used at every stage.

Overview

Existing pretraining-data curation methods define a fixed processing strategy at the corpus or domain level and apply it uniformly to many examples, without adapting to the needs of each example. DataOrchestra is a framework that orchestrates an example-specific pipeline for each chunk of pretraining data. Given a chunk, a small orchestrator model first decides whether to Drop, Untouch, or Clean it. For a chunk to be cleaned, the orchestrator selects one or more downstream stages and, for each rewriting stage, generates a concrete chunk-specific instruction that is executed by the corresponding tool model.

Cleaning is organized into three stages of increasing intervention, applied in the order NP → SR → PA:

• Noise Pruning (NP) — a lightweight small LLM emits whole-line remove_lines(start, end) edits to strip boilerplate, ads, navigation, and other line-level noise.
• Surface Rectification (SR) — an LLM rewrites the chunk to repair broken tables/formulas, grammar, and layout, while preserving the original meaning.
• Pedagogical Augmentation (PA) — an LLM enriches knowledge- or reasoning-intensive text by making reasoning explicit and adding relevant explanation.

The orchestrator is trained by (1) collecting seed documents, (2) having a teacher LLM propose a coarse-grained plan per chunk, (3) evolving that plan by actually running the tool models and verifying each step, which yields a chunk-specific instruction for every retained rewriting stage, and (4) supervised fine-tuning the orchestrator on the resulting pairs.

Table of Contents

• Repository Structure
• Setup
• Quick Start — Run the Orchestrator
• Quick Start — Run the NP Model
• Data Format
• Building the Orchestrator Training Data
  • Stage 1 — Coarse-grained Plan Generation
  • Stage 2 — Fine-grained Plan Evolution
  • Stage 3 — Construct SFT Data
  • Stage 4 — Fine-tune the Orchestrator
• Curating a Corpus with the Trained Orchestrator
• Prompts
• License
• Citation

Repository Structure

inference/
├── common.py                 # chunker (W=1024, Qwen3 tokenizer), async LLM client, JSONL I/O
├── parsing.py                # shared JSON / analysis-result parsing helpers
├── shard_io.py               # split a .jsonl tree into shards / merge shards back
├── file_claim.py             # multi-node O_EXCL claim + heartbeat for distributed runs
├── decision_generation/      # Stage 1: teacher LLM -> coarse plan (drop gate + NP/SR/PA routing)
│   └── prompts/              #   drop.py, decision.py
├── plan_evolution/           # Stage 2: run tool models + verifiers -> fine-grained plan
│   └── prompts/              #   noise_pruning, surface_rectification, pedagogical_augmentation,
│                             #   np_verifier, rewrite_verifier
├── construct_sft/            # Stage 3: build (chunk -> plan) SFT pairs for the orchestrator
└── student/                  # Deployment: run the trained orchestrator + tool models on a corpus
    └── prompts/              #   orchestrator, noise_pruning, surface_rectification, pedagogical_augmentation

scripts/inference/
├── serve_unit.sh             # bring up one vLLM server ("unit") on a GPU subset
├── decision_generation/      # run_distributed.sh / run_node.sh
├── plan_evolution/           # run_distributed.sh + role_config / sampling_overrides examples
└── student/                  # run_distributed.sh + role_config / sampling_overrides examples

Each stage is a self-contained pipeline with two entry points: runner.py (single process, for a quick local test) and distributed_run.py (multi-node; every input .jsonl file is a claim unit written to a mirrored output file, so many workers on many nodes can share the load safely). The run_distributed.sh scripts wrap distributed_run.py, bring up the required vLLM servers, and wait for health.

Setup

The pipelines serve models with vLLM and talk to them over the OpenAI-compatible API.

git clone https://github.com/GAIR-NLP/DataOrchestra
cd DataOrchestra

conda create -n dataorchestra python=3.11
conda activate dataorchestra

pip install vllm                       # serves the orchestrator / tool / teacher models
pip install openai httpx transformers  # async client + chunking tokenizer
pip install rapidfuzz                   # exact Levenshtein for the plan-evolution similarity gate

The framework uses the following models (all served locally through serve_unit.sh); point the scripts at your checkpoints:

Role	Model	Note
Teacher / Verifier	Qwen3-235B-A22B-Instruct-2507	coarse plan + NP/SR/rewrite verification
Orchestrator	fine-tuned Qwen3-1.7B-Base	the model this repo trains and deploys — released at DataOrchestra/Orchestrator
NP tool	fine-tuned Qwen3-0.6B-Base	whole-line noise pruning (ProX-style) — released at DataOrchestra/NP-0.6B
SR / PA tool	Qwen3-4B (non-thinking)	off-the-shelf rewriter

Quick Start — Run the Orchestrator

The orchestrator is released at 🤗 DataOrchestra/Orchestrator. Given one chunk of pretraining text, it emits a compact JSON plan — whether to drop / untouch / clean the chunk, and, when cleaning, which of NP / SR / PA to apply (plus a chunk-specific instruction for each rewriting stage).

The wire format is exactly what the model was fine-tuned on (see inference/student/prompts/orchestrator.py): a one-line system prompt and the raw chunk wrapped in [DOC] / [/DOC].

import json
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL = "DataOrchestra/Orchestrator"
tokenizer = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype="auto", device_map="auto")

# Byte-identical to inference/student/prompts/orchestrator.py
SYSTEM_PROMPT = "You are an excellent orchestrator for pretraining data cleaning."


def plan_for_chunk(chunk: str) -> dict:
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"[DOC]\n{chunk}\n[/DOC]"},  # wrap_doc()
    ]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,           # orchestrator runs non-thinking
    )
    inputs = tokenizer(text, return_tensors="pt").to(model.device)
    generated = model.generate(
        **inputs,
        max_new_tokens=1024,
        do_sample=False,                 # greedy: temperature 0.0 / top_p 1.0
    )
    response = tokenizer.decode(
        generated[0][inputs.input_ids.shape[1]:], skip_special_tokens=True
    )
    return json.loads(response)


chunk = (
    "Home | About | Contact\n\n"
    "The Pythagorean theorem states that a^2 + b^2 = c^2 for a right triangle. "
    "It is one of the most fundamental results in geometry.\n\n"
    "Click here to subscribe to our newsletter!"
)
print(json.dumps(plan_for_chunk(chunk), indent=2, ensure_ascii=False))

The orchestrator returns a flat plan JSON:

{
  "decision": "clean",
  "noise_pruning": true,
  "surface_rectification": "Remove the navigation header and the newsletter call-to-action; keep the statement of the theorem.",
  "pedagogical_augmentation": "Add an intuitive explanation of why a^2 + b^2 = c^2 holds, with a worked example."
}

Plan schema:

Field	Type	Meaning
decision	"drop" | "untouch" | "clean"	top-level gate; only clean triggers the stages below
noise_pruning	bool	run the NP tool model (whole-line remove_lines edits)
surface_rectification	str | null	if a string, run SR with this chunk-specific instruction; null skips
pedagogical_augmentation	str | null	if a string, run PA with this chunk-specific instruction; null skips

The snippet above only produces the plan. To actually execute the plan (apply NP → SR → PA with the tool models and rebuild the cleaned document, with resume, sharding, and multi-node support), use the packaged student driver described in Curating a Corpus with the Trained Orchestrator — it wires the orchestrator together with the NP and SR/PA tool models over vLLM. The parser that turns the raw response into a plan (and gracefully falls back to untouch on malformed output) lives in inference/student/parsers.py::parse_orchestrator_output.

Data Format

Every pipeline reads a directory of .jsonl files (--input_dir), one JSON document per line, with at least an id and a text field:

{"id": "doc-000123", "text": "The full raw document text ...", "meta": {"source": "redpajama-v2"}}

Documents are split into chunks of at most W = 1024 tokens (Qwen3 tokenizer) inside the pipeline — no pre-chunking is needed. If your raw corpus is a few very large .jsonl files, slice it into shards first so workers can share the load at a finer granularity:

python -m inference.shard_io split --input_dir RAW --output_dir SHARDED --shard_size 10000

Building the Orchestrator Training Data

This reproduces the four-stage training-data pipeline shown above. Each run_distributed.sh has its paths and knobs at the top of the file — edit them, then run.

Stage 1 — Coarse-grained Plan Generation

A teacher LLM produces a coarse plan per chunk: first a drop gate (rate the chunk 0–5 under a strict rubric; drop when the score is < 2), then, for each surviving chunk, a binary True/False judgment for each of NP, SR, and PA (all-False ⇒ untouch).

# Edit MODEL_PATH / INPUT_DIR / OUTPUT_DIR at the top of the script, then:
bash scripts/inference/decision_generation/run_distributed.sh

Stage 2 — Fine-grained Plan Evolution

The coarse plan is grounded in actual tool-model execution. For each chunk marked for cleaning, the selected stages run in order and each is checked by a verifier LLM:

• NP is reverted if the verifier finds non-noise lines were over-pruned.
• Each rewriting stage (SR/PA) is verified for noise removal and fidelity (no content loss, no factual error). A failed rewrite is retried under a verifier-written corrective instruction, up to N = 5 times; a stage that keeps failing or barely changes the chunk (Levenshtein similarity > 98%) is dropped. Once a rewrite passes, the stage is annotated with the instruction the verifier synthesizes from that passing before-and-after comparison.

This stage uses two nodes per group: one serves the teacher (which also acts as the three verifiers), the other serves the NP and SR/PA tool models and runs the worker. Pairing is automatic via a shared-filesystem rendezvous directory.

# Edit the *_MODEL_PATH / INPUT_DIR (= Stage 1 output) / OUTPUT_DIR at the top, then
# launch on each node (ROLE defaults to teacher/tools by node rank):
bash scripts/inference/plan_evolution/run_distributed.sh

Per-role endpoints and sampling parameters are configured in scripts/inference/plan_evolution/role_config.example.json and sampling_overrides.example.json.

Stage 3 — Construct SFT Data

Turn the evolved plans into (chunk → plan) supervised fine-tuning pairs. The target is a flat JSON plan: decision plus a boolean for NP and, for SR/PA, either the chunk-specific instruction string or null.

python -m inference.construct_sft.construct_orchestrator_sft \
    /path/to/data/plan_evolution_output \
    --output orchestrator_sft.json

The output is in LLaMA-Factory format (instruction / input / output / system).

Stage 4 — Fine-tune the Orchestrator

Fine-tune Qwen3-1.7B-Base on the pairs from Stage 3 with your SFT trainer of choice (we use LLaMA-Factory). See the paper's appendix for the exact hyper-parameters (3 epochs, LR 3e-5 cosine, global batch size 256, context length 4096).

Curating a Corpus with the Trained Orchestrator

Once you have the orchestrator (and the NP / SR-PA tool models), run the full example-specific pipeline over a raw corpus. The student driver reads raw documents, chunks them, runs the orchestrator to produce a plan per chunk, executes the selected tool models, and rebuilds the cleaned document.

# Edit ORCH_MODEL_PATH / NP_MODEL_PATH / SR_PA_MODEL_PATH and INPUT_DIR / OUTPUT_DIR, then:
bash scripts/inference/student/run_distributed.sh

The driver defaults to PHASE=chain, running orchestrator → NP → SR/PA back to back and swapping the model only when a phase is globally complete. You can also run a single PHASE (orchestrator / np / sr_pa) at a time. To export the cleaned corpus as flat training data, merge the per-document output and keep only the fields you need:

python -m inference.shard_io merge --input_dir OUTPUT_DIR --output_dir CLEAN_CORPUS

Two useful ablation switches are exposed on the student driver: --ignore_orchestrator_instructions (keep the stage on/off decisions but drop the chunk-specific instructions) and --stage (re-run from a specific stage).

Prompts

All prompts are kept in code, one module per role, and mirror the figures in the paper's appendix:

• Coarse plan: inference/decision_generation/prompts/{drop,decision}.py
• Tool models: inference/plan_evolution/prompts/{noise_pruning,surface_rectification,pedagogical_augmentation}.py
• Verifiers: inference/plan_evolution/prompts/{np_verifier,rewrite_verifier}.py
• Deployed orchestrator + tools: inference/student/prompts/*.py

The NP tool prompt and the orchestrator prompt are kept byte-identical between the training-data construction side and the deployment side, so the released models stay in-distribution.

License

This repository is released under the Apache License 2.0.

Citation

If you find this work useful, please cite:

@article{dataorchestra2026,
  title   = {DataOrchestra: Learning to Orchestrate Per-Example Curation of Pretraining Data},
  author  = {Huang, Zhen and Wang, Yikun and Xia, Shijie and Liu, Pengfei},
  year    = {2026},
  journal = {arXiv preprint arXiv:2607.24717}
}