DataOrchestra — Orchestrator Model

Model Details

Base model	Qwen/Qwen3-1.7B-Base
Role	Orchestrator (plan generator)
Input	one pretraining-data chunk (≤ 1024 Qwen3 tokens)
Output	a flat JSON plan (decision + NP / SR / PA)
Inference mode	non-thinking, greedy decoding

Usage

The wire format is a one-line system prompt plus the raw chunk wrapped in [DOC] / [/DOC]. The model responds with a single JSON plan.

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")

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]"},
    ]
    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))

Serving with vLLM

For high-throughput curation, serve the model with an OpenAI-compatible endpoint:

vllm serve DataOrchestra/Orchestrator --served-model-name DataOrchestra-Orchestrator --trust-remote-code

from openai import OpenAI

client = OpenAI(base_url="http://127.0.0.1:8000/v1", api_key="EMPTY")
resp = client.chat.completions.create(
    model="DataOrchestra-Orchestrator",
    messages=[
        {"role": "system", "content": "You are an excellent orchestrator for pretraining data cleaning."},
        {"role": "user", "content": "[DOC]\n<your chunk here>\n[/DOC]"},
    ],
    temperature=0.0,
    max_tokens=1024,
    extra_body={"chat_template_kwargs": {"enable_thinking": False}},
)
print(resp.choices[0].message.content)

Output Schema

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."
}

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

For drop / untouch decisions, all three stage fields are inert.

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}
}