TACO: Agentic SFT Data Curation from Multi-distributional Corpora via User Instructions

Code release for the paper Agentic SFT Data Curation from Multi-distributional Corpora via User Instructions (anonymous submission).

TACO is a four-stage agentic pipeline for instruction-tuning data curation. Each stage is an LLM agent that inspects the data pool, reads the user's natural-language task instruction $\tau$, infers a per-run configuration, and applies it. The four stages decouple the four mixture axes of real instruction pools one axis at a time:

Stage	Axis it decouples	Output configuration
Stage 0 — Workflow Mode Detection	Domain	matched cluster IDs
Stage 1 — Rule-Based Pre-Filter	Surface characteristics	content-quality policy
Stage 2 — Task-Adaptive Statistical Filter	Informativeness (relative to base model)	four-metric weights
Stage 3 — Task-Adaptive Fine Filter	Quality (absolute)	open-vocabulary judging rubric

---

Setup

The project targets Python 3.12 and uses uv for dependency management. Full pipeline execution requires CUDA 12.x and at least one A100-80GB-class GPU; the agent and preprocessing stages run on CPU.

git clone https://github.com/nju-websoft/TACO.git
cd TACO
uv sync
source .venv/bin/activate

Dependencies are pinned in uv.lock; the high-level groups are langgraph + langchain-openai (agent runtime and orchestrator LLM client), sentence-transformers + hdbscan + faiss-cpu (Stage 0 clustering and labeling), vllm (base-model inference for Stage 2 metrics and downstream evaluation), and deepspeed + accelerate (full-parameter SFT under ZeRO-3).

Configuration

Copy the template and fill in the marked fields:

cp config.yaml.example config.yaml

config.yaml is in .gitignore; the real values never leave your machine. The minimum you must set:

bd:
  # Base model used in Stage 2 metric computation and downstream SFT
  base_model_path: /path/to/Qwen2.5-7B           # local HF snapshot
  base_model_url:  http://localhost:8250/v1      # local vLLM endpoint

  # Sentence encoder used by Stage 0
  embed_model_path: BAAI/bge-m3                  # HF identifier or local path
  embed_url:        http://localhost:8003/v1     # local vLLM endpoint

  # Frozen orchestrator/judge LLM (OpenAI-compatible API)
  base_url: https://api.openai.com/v1
  api_key:  ${OPENAI_API_KEY}
  model:    gpt-4o-mini                          # used by all stages
  subagent_model: gpt-4o-mini

target_num: 10000                                # final curated subset size
rough_filter_min_retention: 0.85
fine_filter:
  samples_per_round: 50
  num_batches: 5

You will also need two vLLM servers running locally before launching the pipeline, one for the base model and one for the sentence encoder. The endpoints are pointed at by base_model_url and embed_url.

Running the pipeline

End-to-end (preprocess + agent curation + SFT):

python run_pipeline.py \
    --instruction "Extract high-quality math reasoning samples from this multi-domain instruction pool." \
    --origin_data_dir /path/to/your/pool \
    --model_path /path/to/Qwen2.5-7B \
    --gpus 0,1,2,3 \
    --epochs 3 \
    --deepspeed_config ds_config/ds_config_7b.json \
    --test_data_dirs /path/to/benchmark/gsm8k /path/to/benchmark/math

The driver runs three sequential stages.

1. Preprocess (preprocess/preprocess.py) — adds stable per-sample IDs and bge-m3 embeddings to the input pool. Idempotent; the --resume flag skips this when re-running.
2. Agent (agent_interface.chat) — executes the four TACO stages on the preprocessed pool and writes a <pool>_<model>_rough_basemodel_fine/ subdirectory containing the curated 10k subset and the per-stage artefact files (cluster labels, rule policy JSON, metric weights, judging rubric).
3. Train (finetune/train.py) — full-parameter SFT on the curated subset under DeepSpeed ZeRO-3, with per-epoch evaluation on the --test_data_dirs benchmarks.

For Stage 0 only (clustering + labeling + visualisation), use run_stage0_viz.sh; outputs land in stage0_runs/<timestamp>/ and include the cluster embeddings, free-text labels, the matched cluster set, and a UMAP projection.

The full list of flags is in python run_pipeline.py --help.

Repository layout

.
├── agent/                     # TACO agent runtime
│   ├── components/            # Planner, executor, graph nodes
│   ├── tools/
│   │   ├── filters/           # rough_filter / base_model_filter / fine_filter / discipline_filter
│   │   ├── subagents/         # one sub-agent per stage
│   │   └── subagents_prompt/  # one prompt template per sub-agent
│   ├── dispatch.py            # tool-call dispatcher
│   ├── graph.py               # LangGraph state machine
│   └── workflow_modes.py      # multi-domain vs single-domain mode logic
├── preprocess/                # ID + embedding pre-pass
├── finetune/                  # Full-parameter SFT + per-benchmark evaluator harness
│   └── test/                  # math / code / MCQ / LiveCodeBench evaluators
├── ds_config/                 # DeepSpeed ZeRO-3 configs (0.5B / 1.5B / 3B / 7B)
├── utils/                     # Data formatting and ancillary scripts
├── agent_interface.py         # Entry point for the agent pipeline
├── run_pipeline.py            # End-to-end Preprocess → Agent → Train driver
├── run_stage0_viz.py          # Stage 0 cluster-and-label visualiser
├── compress_dataset.sh        # Pack a curated subset for release
└── config.yaml.example        # Configuration template

Datasets and benchmarks

The training pools used in the paper (OpenHermes-2.5, Tulu-3-SFT-Mixture, OpenR1-Math, Synthetic-Math, Medical-Reasoning, UltraMedical, OSS-Instruct, Code-Alpaca) are public Hugging Face datasets and are not redistributed with this repository.

The evaluation harness consumes 11 standard benchmarks (GSM8K, MATH, Minerva-Math, AMC23, AIME 2024/2025, HumanEval, LiveCodeBench, BigCodeBench, MedQA, MMLU-Medical) plus three science benchmarks (MMLU-STEM, GPQA-Diamond, SciBench) used in §4.5. Each benchmark is loaded from a directory of JSON/JSONL/Parquet files; the directory layout expected by finetune/test/ is documented in finetune/test/README.md.

Reproducing the paper

The five baseline methods compared in §4.2 (Random, SuperFiltering, Select2Reason, MIG, SelectIT) and the two two-step baselines in §4.5 (data_source-tag and seed-embedding retrieval, both feeding into SuperFiltering) are implemented per their respective original papers and code releases; the implementations are not included here.

For TACO itself, the per-stage artefact directories written by agent_interface.chat contain everything needed to inspect a particular run, including the LLM-emitted configuration JSONs at every stage, the surviving sample IDs after each filter, and the final per-sample scores.

Citation

@inproceedings{taco2026,
  title  = {Agentic SFT Data Curation from Multi-distributional Corpora via User Instructions},
  author = {Anonymous},
  year   = {2026},
  note   = {Under review.}
}

License

Apache License 2.0 (see LICENSE).