DataMaster

Official implementation of DataMaster, the method introduced in Agentic Instruction Data Selection: Let DataMaster Interpret Your Intent.

DataMaster turns a natural-language description of the desired training data into an executable selection process. Instead of applying one fixed scoring rule to every task, it inspects the source SFT pool and coordinates up to four specialized agents:

1. Domain Agent identifies instruction-aligned domains in a mixed pool.
2. Characteristic Agent derives task-dependent filtering rules from the pool schema, statistics, and content patterns.
3. Informativeness Agent combines target-model NLL, token-level entropy, semantic drift, and MeanDiff using instruction-conditioned weights.
4. Quality Agent constructs a task-specific rubric, evaluates candidates, and retains the requested budget.

The Domain Agent is used for multi-domain selection. Single-domain selection starts from the Characteristic Agent. The arXiv link and citation entry will be added after the public paper identifier is available.

Repository contents

.
|-- agent/                  # orchestrator, stage agents, prompts, and tools
|-- preprocess/             # IDs, profiles, embeddings, and FAISS indexing
|-- finetune/               # full-parameter SFT and evaluation harness
|-- baselines/              # five paper-aligned baseline runners
|-- ds_config/              # three ZeRO-3 learning-rate configurations
|-- examples/toy_pool/      # tiny synthetic input example
|-- docs/DATASETS.md        # source-pool and benchmark preparation
|-- run_selection.py        # preprocessing plus DataMaster selection
|-- run_pipeline.py         # selection, SFT, and evaluation
|-- run_stage0_viz.py       # optional Domain-Agent visualization
|-- config.yaml.example     # configuration template
|-- SECURITY.md             # generated-code execution warning
`-- THIRD_PARTY_NOTICES.md  # incorporated and adapted third-party code

Model weights, full source pools, selected subsets, embeddings, benchmarks, checkpoints, logs, and API traces are not stored in this repository. Public source identifiers and the regeneration procedure are documented in docs/DATASETS.md.

Paper-to-code map

Paper component	Agent/controller	Core implementation
Workflow routing	agent/graph.py, agent/components/planner.py	compiled agent graph and workflow-mode detection
Stage 1: Domain Agent	agent/tools/subagents/discipline_discovery_agent.py	run_discipline_discovery in discipline_filter.py
Stage 2: Characteristic Agent	agent/tools/subagents/rough_filter_agent.py	run_rough_filter in rough_filter.py
Stage 3: Informativeness Agent	agent/tools/subagents/base_model_filter_agent.py	BaseModelTools in base_model_filter.py
Stage 4: Quality Agent	agent/tools/subagents/fine_filter_agent.py	run_fine_filter_single_file in fine_filter.py
Full-parameter SFT	finetune/train.py	DeepSpeed ZeRO-3 training and benchmark dispatch
Evaluation	finetune/test/	math, medical, science, and code evaluators

The Domain Agent uses HDBSCAN when the current pool contains at most 50,000 samples. For larger pools, or when HDBSCAN does not produce a usable partition, the implementation falls back to KMeans.

Requirements

• Python 3.12
• Linux is recommended for vLLM, DeepSpeed, and multi-GPU training
• CUDA-capable hardware for target-model signals, SFT, and evaluation
• an OpenAI-compatible endpoint for the agent LLM
• an OpenAI-compatible embedding endpoint and local embedding tokenizer/model
• local or Hugging Face target-model weights

Create an environment with uv:

uv sync
source .venv/bin/activate

or install with pip:

python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install -r requirements.txt

The dependency set includes training and all evaluation backends. A selection-only environment may omit DeepSpeed and benchmark-specific packages that are not used in the intended run.

Configuration

Copy the template and fill in local model identifiers and endpoints:

cp config.yaml.example config.yaml

Important fields are:

Field	Meaning
bd.base_api_model_name	target-model identifier exposed by the local endpoint
bd.base_model_path	local target-model/tokenizer path
bd.base_model_url	OpenAI-compatible target-model endpoint
bd.embed_api_model_name	embedding model identifier
bd.embed_model_path	local embedding tokenizer/model path; BGE-M3 is used by the paper setup
bd.embed_url	OpenAI-compatible embedding endpoint
bd.base_url	agent-LLM endpoint
bd.api_key	agent-LLM credential; keep this only in untracked config.yaml
bd.model, bd.subagent_model	orchestrator and stage-agent model identifiers
target_num	final selection budget; the paper uses 10,000
enable_rough_filter_retention_gate	optional engineering safeguard, disabled for paper reproduction

The bd key is the configuration namespace expected by the released code. config.yaml is ignored by Git. Do not commit API keys, private endpoints, or local filesystem paths.

Input data

Create one directory per SFT pool. Each top-level JSON file must contain a list of records:

dataset/
`-- my_pool/
    |-- data_000.json
    `-- data_001.json

Each record should provide:

{
  "instruction": "Solve the problem and explain the reasoning.",
  "input": "If x + 2 = 5, what is x?",
  "output": "x = 3.",
  "text": "Below is an instruction ... ### Response:\\nx = 3."
}

The input field may be empty. Keep generated filtered_*.json, embeddings, indexes, and outputs outside a fresh source directory.

Dataset-card profile

DataMaster combines deterministic statistics from the exact local split with a normalized summary of the upstream Dataset Card. Start from examples/dataset_card_profile.example.yaml and record the dataset ID, revision, split, Dataset Card URL, domain coverage, languages, construction, and intended format. Do not guess an unknown revision.

Passing the summary to preprocessing creates profile.txt in the pool directory. That generated profile is the only profile consumed by the agent.

Run selection

The smallest complete command is:

python run_selection.py \
  --instruction "Select mathematical reasoning data with rigorous, verifiable solutions." \
  --data-dir dataset/my_pool \
  --dataset-card-profile examples/dataset_card_profile.example.yaml \
  --refresh-profile

Preprocessing assigns stable working-copy IDs, computes embeddings, builds the FAISS index, and writes the pool profile. DataMaster then selects records using the workflow implied by the instruction and inspected pool.

For an already preprocessed pool:

python run_selection.py \
  --instruction "Select mathematical reasoning data with rigorous, verifiable solutions." \
  --data-dir dataset/my_pool \
  --skip-preprocess

Use --resume only for a compatible interrupted run. The selector uses seed 42 by default; pass --seed to change it.

For multi-domain selection, explicitly state that the pool is multi-domain and name the target domain, for example:

python run_selection.py \
  --instruction "From this multi-domain SFT pool, select code data that emphasizes executable solutions and instruction compliance." \
  --data-dir dataset/mixed_pool \
  --dataset-card-profile examples/dataset_card_profile.example.yaml \
  --refresh-profile

Domain-Agent visualization

The optional visualization entry point runs domain discovery independently:

python run_stage0_viz.py \
  --data dataset/mixed_pool \
  --instruction "Select medical reasoning data" \
  --output outputs/domain_visualization

Run python run_stage0_viz.py --help for clustering and sampling options.

Fine-tuning and evaluation

The repository contains three DeepSpeed ZeRO-3 configurations that differ only in learning rate:

ds_config/ds_config_lr_1e-6.json
ds_config/ds_config_lr_5e-6.json
ds_config/ds_config_lr_1e-5.json

Train a fixed selected subset with:

deepspeed --include localhost:0,1 finetune/train.py \
  --data_dir /path/to/selected_subset \
  --original_data_dir /path/to/source_pool \
  --model_path /path/to/target_model \
  --epochs 3 \
  --test_data_dirs /path/to/benchmark_a /path/to/benchmark_b \
  --deepspeed \
  --deepspeed_config ds_config/ds_config_lr_5e-6.json

The optional training-time Judger is disabled by default so the selected subset remains fixed across epochs. --enable_judger enables that separate engineering extension.

The end-to-end entry point chains preprocessing, selection, SFT, and evaluation:

python run_pipeline.py \
  --instruction "Select mathematical reasoning data with rigorous, verifiable solutions." \
  --model_path /path/to/target_model \
  --origin_data_dir dataset/my_pool \
  --dataset_card_profile examples/dataset_card_profile.example.yaml \
  --refresh_profile \
  --test_data_dirs /path/to/gsm8k /path/to/math \
  --deepspeed_config ds_config/ds_config_lr_5e-6.json \
  --gpus 0,1

Benchmark files are not redistributed. Preparation and expected directory names are described in docs/DATASETS.md and finetune/test/README.md.

Reproducing the experiment families

The paper evaluates:

• six single-domain SFT pools covering math, medicine, and code;
• two mixed-domain SFT pools with math and code targets, plus a science target on OpenHermes-2.5;
• Qwen2.5-7B, Llama-3.1-8B, and Olmo-3-7B-Instruct target models;
• a fixed 10K selection budget for selection methods;
• three candidate learning rates selected before the large-scale comparison;
• three training epochs, with reported scores aggregated over the final third epoch according to the paper protocol.

To reproduce a setting, obtain the official source pool, normalize it to the input schema, create a revision-pinned Dataset Card summary, run DataMaster with the corresponding natural-language instruction, train the target model with the selected fixed learning rate, and evaluate on the domain benchmarks. This repository does not claim byte-identical reconstruction when upstream dataset revisions or hosted model services change.

Baselines

baselines/ provides paper-aligned runners for Random, SuperFiltering, SelectIT, DSIR, and TSDS. DSIR and TSDS retain their official MIT-licensed cores. See baselines/README.md for input contracts and commands. Baseline-selected subsets and result tables are not included.

Security

Code-generation evaluation executes model-produced Python. Run it only inside a disposable, least-privilege sandbox without credentials, personal mounts, or unrestricted network access. Evaluator timeouts are reliability controls, not a security boundary. Read SECURITY.md before running code benchmarks.

Data and licensing

Project-authored code is released under Apache License 2.0. Vendored or adapted third-party components retain their original licenses and notices; see THIRD_PARTY_NOTICES.md and third_party_licenses/.

Dataset records are not covered by the repository's code license and are not redistributed here. Users are responsible for obtaining datasets from their official sources and complying with their licenses and terms.

Citation

The citation entry will be added after the arXiv identifier is available.