📜 Paper | 🤗 Datasets | 🤗 Models | 📦 UltraData Collection
English | 中文
UltraX is a function-calling refinement framework for large-scale pre-training data that adaptively generates and executes editing functions for efficient instance-wise refinement. Instead of relying on end-to-end text rewriting, UltraX trains a lightweight refinement model to predict structured editing operations — including insertion, deletion, and modification — which are then deterministically executed on the original text.
Key Features:
add_line), deletion (remove_lines, remove_all), and modification (replace_str), enabling fine-grained instance-level editing beyond delete-only approaches.keep_all, remove_all, remove_lines, replace_str, add_line), enabling fine-grained instance-level editing with deterministic execution.replace_str operations with unique context anchoring.The UltraX pipeline consists of two main stages:
| Step | Module | Description |
|---|---|---|
| 1 | seed_preprocessing/ | Sliding window splitter for long documents (12K token windows, 20% overlap) |
| 2 | prompt_optimization/ | Dataset-adaptive prompt optimization with profiling, iterative refinement, and regression testing |
| 3 | e2e_refinement/ | Batch end-to-end text refinement via expert LLM API |
| 4 | function_construction/ | LAM + DCR: convert (original, refined) pairs into structured function call training data |
| 5 | sft_data_building/ | Edit-weighted sampling + system instruction injection |
| 6 | model_training/ | Full-parameter SFT with ms-swift + DeepSpeed ZeRO3 |
| Step | Module | Description |
|---|---|---|
| 7 | inference/ | Multi-GPU data-parallel vLLM inference with sliding window segmentation |
| 8 | post_processing/ | Function parsing, validation, post-processing, and deterministic execution |
We pretrain 1B-parameter MiniCPM models from scratch on 20B tokens and evaluate on 10 benchmarks with zero-shot setting.
UltraX achieves the highest average performance on all five corpora, winning 34 out of 50 task-corpus pairs.
Average downstream performance on FineWeb under different training token budgets.
git clone https://github.com/BIGWangYuDong/UltraX.git
cd UltraX
conda create -n ultrax python=3.10
conda activate ultrax
pip install -r requirements.txt
For inference, you also need vLLM:
pip install vllm
For model training, install ms-swift:
pip install ms-swift
Split long documents into overlapping windows at line boundaries:
python stage1_model_construction/seed_preprocessing/sliding_window_splitter.py \
--input_dir /path/to/seed_data \
--output_dir /path/to/output \
--tokenizer_path /path/to/tokenizer \
--max_tokens 12000 \
--overlap_ratio 0.2
Automatically optimize refinement prompts for each dataset:
cd stage1_model_construction/prompt_optimization
python main.py \
--api-key $API_KEY \
--datasets fineweb redpajama \
--max-iterations 200 \
--batch-size 5
Refine seed data using optimized prompts via expert LLM:
python stage1_model_construction/e2e_refinement/refine_dataset.py \
--input_dir /path/to/seed_data \
--output_dir /path/to/refined_output \
--prompt_dir /path/to/optimized_prompts \
--api_url $API_URL \
--api_key $API_KEY \
--model deepseek-v3
Convert (original, refined) text pairs into structured function call training data:
python stage1_model_construction/function_construction/function_construction.py \
--input_dir /path/to/refined_data \
--output_dir /path/to/training_data \
--num_workers 64
Sample training data by operation combination and add system instructions:
python stage1_model_construction/sft_data_building/sample_and_format.py \
--train_data_dir /path/to/training_data \
--output_dir /path/to/sft_data
Train the refinement model with ms-swift:
bash stage1_model_construction/model_training/train.sh
Note: Edit the paths in
train.shbefore running. See the script for detailed hyperparameter settings.
Run multi-GPU data-parallel inference with vLLM:
python stage2_large_scale_execution/inference/inference.py \
--input_dir /path/to/raw_data \
--model_path /path/to/trained_model \
--output_dir /path/to/inference_output \
--num_gpus 8 \
--max_chars 48000
Parse, validate, and execute the predicted cleaning operations:
python stage2_large_scale_execution/post_processing/post_process_and_execute.py \
--input_dir /path/to/inference_output \
--output_dir /path/to/cleaned_data \
--num_workers 128
The output contains 3 columns: original, cleaned, processed_functions.
| Function | Description |
|---|---|
keep_all() | Document needs no modification |
remove_all() | Entire document is valueless (e.g., error pages, login walls) |
remove_lines(start, end) | Remove consecutive lines from start to end (inclusive) |
replace_str(line, old, new) | Replace a substring within a specific line |
add_line(base, sub_idx, content) | Insert a new line near the base position |
We thank the following projects for their great contributions:
If you find our work useful, please consider citing:
@misc{ultrax2026,
title={UltraX: Refining Pre-Training Data at Scale with Adaptive Programmatic Editing},
author={Xinlong Zhao and Dongsheng Liu and Hengyu Zhao and Zixuan Fu and Zheng Wang and Jie Cai and Jie Zhou and Qiang Ma and Xuanhe Zhou and Xu Han and Yudong Wang and Zhiyuan Liu},
year={2026},
eprint={2607.08646},
archivePrefix={arXiv},
primaryClass={cs.CL},
}
This project is licensed under the Apache 2.0 license.
No unauthorized unchanged redistribution: Without prior written permission from the original authors (or this organization), any institution, organization, or third-party platform is strictly prohibited from directly reposting, mirroring, re-hosting, or commercially repackaging and republishing any artifacts of this project in any form.
📜 Paper | 🤗 Datasets | 🤗 Models | 📦 UltraData Collection
English | 中文
UltraX is a function-calling refinement framework for large-scale pre-training data that adaptively generates and executes editing functions for efficient instance-wise refinement. Instead of relying on end-to-end text rewriting, UltraX trains a lightweight refinement model to predict structured editing operations — including insertion, deletion, and modification — which are then deterministically executed on the original text.
Key Features:
add_line), deletion (remove_lines, remove_all), and modification (replace_str), enabling fine-grained instance-level editing beyond delete-only approaches.keep_all, remove_all, remove_lines, replace_str, add_line), enabling fine-grained instance-level editing with deterministic execution.replace_str operations with unique context anchoring.The UltraX pipeline consists of two main stages:
| Step | Module | Description |
|---|---|---|
| 1 | seed_preprocessing/ | Sliding window splitter for long documents (12K token windows, 20% overlap) |
| 2 | prompt_optimization/ | Dataset-adaptive prompt optimization with profiling, iterative refinement, and regression testing |
| 3 | e2e_refinement/ | Batch end-to-end text refinement via expert LLM API |
| 4 | function_construction/ | LAM + DCR: convert (original, refined) pairs into structured function call training data |
| 5 | sft_data_building/ | Edit-weighted sampling + system instruction injection |
| 6 | model_training/ | Full-parameter SFT with ms-swift + DeepSpeed ZeRO3 |
| Step | Module | Description |
|---|---|---|
| 7 | inference/ | Multi-GPU data-parallel vLLM inference with sliding window segmentation |
| 8 | post_processing/ | Function parsing, validation, post-processing, and deterministic execution |
We pretrain 1B-parameter MiniCPM models from scratch on 20B tokens and evaluate on 10 benchmarks with zero-shot setting.
UltraX achieves the highest average performance on all five corpora, winning 34 out of 50 task-corpus pairs.
Average downstream performance on FineWeb under different training token budgets.
git clone https://github.com/BIGWangYuDong/UltraX.git
cd UltraX
conda create -n ultrax python=3.10
conda activate ultrax
pip install -r requirements.txt
For inference, you also need vLLM:
pip install vllm
For model training, install ms-swift:
pip install ms-swift
Split long documents into overlapping windows at line boundaries:
python stage1_model_construction/seed_preprocessing/sliding_window_splitter.py \
--input_dir /path/to/seed_data \
--output_dir /path/to/output \
--tokenizer_path /path/to/tokenizer \
--max_tokens 12000 \
--overlap_ratio 0.2
Automatically optimize refinement prompts for each dataset:
cd stage1_model_construction/prompt_optimization
python main.py \
--api-key $API_KEY \
--datasets fineweb redpajama \
--max-iterations 200 \
--batch-size 5
Refine seed data using optimized prompts via expert LLM:
python stage1_model_construction/e2e_refinement/refine_dataset.py \
--input_dir /path/to/seed_data \
--output_dir /path/to/refined_output \
--prompt_dir /path/to/optimized_prompts \
--api_url $API_URL \
--api_key $API_KEY \
--model deepseek-v3
Convert (original, refined) text pairs into structured function call training data:
python stage1_model_construction/function_construction/function_construction.py \
--input_dir /path/to/refined_data \
--output_dir /path/to/training_data \
--num_workers 64
Sample training data by operation combination and add system instructions:
python stage1_model_construction/sft_data_building/sample_and_format.py \
--train_data_dir /path/to/training_data \
--output_dir /path/to/sft_data
Train the refinement model with ms-swift:
bash stage1_model_construction/model_training/train.sh
Note: Edit the paths in
train.shbefore running. See the script for detailed hyperparameter settings.
Run multi-GPU data-parallel inference with vLLM:
python stage2_large_scale_execution/inference/inference.py \
--input_dir /path/to/raw_data \
--model_path /path/to/trained_model \
--output_dir /path/to/inference_output \
--num_gpus 8 \
--max_chars 48000
Parse, validate, and execute the predicted cleaning operations:
python stage2_large_scale_execution/post_processing/post_process_and_execute.py \
--input_dir /path/to/inference_output \
--output_dir /path/to/cleaned_data \
--num_workers 128
The output contains 3 columns: original, cleaned, processed_functions.
| Function | Description |
|---|---|
keep_all() | Document needs no modification |
remove_all() | Entire document is valueless (e.g., error pages, login walls) |
remove_lines(start, end) | Remove consecutive lines from start to end (inclusive) |
replace_str(line, old, new) | Replace a substring within a specific line |
add_line(base, sub_idx, content) | Insert a new line near the base position |
We thank the following projects for their great contributions:
If you find our work useful, please consider citing:
@misc{ultrax2026,
title={UltraX: Refining Pre-Training Data at Scale with Adaptive Programmatic Editing},
author={Xinlong Zhao and Dongsheng Liu and Hengyu Zhao and Zixuan Fu and Zheng Wang and Jie Cai and Jie Zhou and Qiang Ma and Xuanhe Zhou and Xu Han and Yudong Wang and Zhiyuan Liu},
year={2026},
eprint={2607.08646},
archivePrefix={arXiv},
primaryClass={cs.CL},
}
This project is licensed under the Apache 2.0 license.
No unauthorized unchanged redistribution: Without prior written permission from the original authors (or this organization), any institution, organization, or third-party platform is strictly prohibited from directly reposting, mirroring, re-hosting, or commercially repackaging and republishing any artifacts of this project in any form.