OpenBMB/UltraX

OpenBMB/UltraX

Language: Python

License: Apache-2.0

Stars: 2

Forks: 0

Open issues: 0

Created: 2026-06-10T09:27:20Z

Pushed: 2026-07-12T14:18:40Z

Default branch: main

Fork: no

Archived: no

README:

UltraX: Refining Pre-Training Data at Scale with Adaptive Programmatic Editing

📜 Paper | 🤗 Datasets | 🤗 Models | 📦 UltraData Collection

English | 中文

📚 Introduction

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:

• Complete Function Space: Covers insertion (add_line), deletion (remove_lines, remove_all), and modification (replace_str), enabling fine-grained instance-level editing beyond delete-only approaches.
• Reliable Seed Supervision: Uses expert LLM end-to-end refinement + LAM (Line Alignment and Mapping) + DCR (Dynamic Context Replacement) to produce high-quality structured program supervision.
• Robust Large-Scale Execution: Sliding-window inference with overlap-aware aggregation, ambiguity filtering, same-line operation merging, and duplicate pattern detection ensure reliable execution at scale.
• State-of-the-Art Performance: UltraX achieves the best average performance across all five corpora in 1B-model pre-training experiments, with improvements exceeding 2% on multiple datasets.

📢 News

• [2026.07.13] UltraX codebase, refinement model, and refined datasets are now available on GitHub and Hugging Face. 🚀🚀🚀
• [2025.07.10] UltraX technical report is available on arXiv. 🔥🔥🔥

💡 Highlights

• Function-Calling Refinement: Instead of end-to-end text rewriting, UltraX predicts structured editing operations (keep_all, remove_all, remove_lines, replace_str, add_line), enabling fine-grained instance-level editing with deterministic execution.
• LAM + DCR Pipeline: Line Alignment and Mapping (LAM) aligns original and refined text at line level, while Dynamic Context Replacement (DCR) converts character-level edits into reliable replace_str operations with unique context anchoring.
• Robust Large-Scale Execution: Sliding-window inference with overlap-aware aggregation, ambiguity filtering, same-line operation merging, and duplicate pattern detection ensure reliable execution at scale.
• State-of-the-Art Performance: UltraX achieves the best average performance across all five corpora in 1B-model pre-training experiments, with improvements exceeding 2% on multiple datasets and superior data efficiency.

📋 Pipeline Overview

The UltraX pipeline consists of two main stages:

Stage I: Refinement Model Construction

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

Stage II: Large-Scale Program Execution

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

📈 Evaluation Results

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.

🚀 Quick Start

Setup

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

Stage I: Refinement Model Construction

Step 1: Seed Data Preprocessing

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

Step 2: Dataset-Adaptive Prompt Optimization

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

Step 3: End-to-End Refinement

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

Step 4: Function Construction (LAM + DCR)

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

Step 5: SFT Data Building

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

Step 6: Model Training

Train the refinement model with ms-swift:

bash stage1_model_construction/model_training/train.sh

> Note: Edit the paths in train.sh before running. See the script for detailed hyperparameter settings.

Stage...