Orchestrator / README.md
Benjamin02's picture
Import from GAIR/DataOrchestra; rename to Orchestrator, update model paths and arXiv link
d461963 verified
|
Raw
History Blame Contribute Delete
4.19 kB
---
license: apache-2.0
base_model: Qwen/Qwen3-1.7B-Base
language:
- en
library_name: transformers
pipeline_tag: text-generation
tags:
- pretraining-data
- data-curation
- data-cleaning
- dataorchestra
---
# DataOrchestra — Orchestrator Model
## Model Details
| | |
| --- | --- |
| Base model | [`Qwen/Qwen3-1.7B-Base`](https://huggingface.co/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.
```python
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:
```bash
vllm serve DataOrchestra/Orchestrator --served-model-name DataOrchestra-Orchestrator --trust-remote-code
```
```python
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:
```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:
```bibtex
@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}
}
```