📄 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.
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.
- Repository Structure
- Setup
- Quick Start — Run the Orchestrator
- Quick Start — Run the NP Model
- Data Format
- Building the Orchestrator Training Data
- Curating a Corpus with the Trained Orchestrator
- Prompts
- License
- Citation
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.
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 gateThe 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 |
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.
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 10000This 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.
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.shThe 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 = 5times; 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.shPer-role endpoints and sampling parameters are configured in
scripts/inference/plan_evolution/role_config.example.json and sampling_overrides.example.json.
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.jsonThe output is in LLaMA-Factory format
(instruction / input / output / system).
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).
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.shThe 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_CORPUSTwo 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).
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.
This repository is released under the Apache License 2.0.
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}
}