The AITraining interactive wizard provides a modern, user-friendly way to configure and launch training jobs for all supported trainer types—LLMs, text/vision/tabular models, seq2seq, extractive QA, sentence transformers, VLM, and more—from the command line. It guides you through all necessary steps with helpful prompts, validation, and confirmation.
- Step-by-step configuration: Clear progression through trainer selection, dataset setup, model choice, and advanced parameters
- Trainer-specific guidance: Automatically adapts questions based on selected trainer (LLM SFT/DPO/ORPO/PPO, text classification, token classification, tabular, image classification/regression, seq2seq, extractive QA, sentence transformers, VLM)
- Smart defaults: Press Enter to accept sensible defaults shown in brackets
- Validation: Real-time Pydantic validation with helpful error messages
- Configuration summary: Review all settings before starting training
- Graceful cancellation: Press Ctrl+C at any time to exit
- Live logs by default:
log=wandb+ LEET visualizer are now enabled automatically for CLI/TUI workflows
There are three ways to launch the interactive wizard:
Simply run aitraining without any arguments:
aitrainingThis automatically launches the wizard with a friendly interface.
From the LLM subcommand:
aitraining llm --interactiveThe wizard launches automatically when you specify --train but omit required parameters:
# This will launch the wizard since model, data, and project are missing
aitraining llm --train
# This will also launch the wizard to fill in missing details
aitraining llm --train --model gpt2Before Step 0, the wizard now checks for a Hugging Face token so it can list private models/datasets and authenticate any Hub calls:
- If
HF_TOKENis already exported (or you passed--token), the wizard reuses it and surfaces a ✅ confirmation. - Otherwise it prompts for a token (press Enter to skip and continue with public assets only).
- The captured token is exported to
HF_TOKENfor the rest of the session and is still editable later under Advanced → Hub Integration.
Skip the prompt only if you are sure every asset you need is public.
When launched from the base aitraining command (no subcommand), the wizard first prompts you to choose which trainer family you want to configure:
📋 Step 0: Choose Trainer Type
============================================================
Available trainer types:
1 . Large Language Models (LLM) - text generation, chat, instruction following
2 . Text Classification - categorize text into labels
3 . Token Classification - NER, POS tagging
4 . Tabular Data - classification or regression on structured data
5 . Image Classification - categorize images
6 . Image Regression - predict values from images
7 . Sequence-to-Sequence - translation, summarization
8 . Extractive QA - answer questions from context
9 . Sentence Transformers - semantic similarity embeddings
10. Vision-Language Models - multimodal tasks
Select trainer type [1-10, default: 1]:
Once trainer type is selected, the wizard switches into that trainer’s flow. For LLMs, you’ll see an additional sub-step to pick the LLM trainer variant.
For LLM workflows, the wizard then prompts you to choose the specific LLM trainer variant:
📋 Step 1: Choose Training Type
============================================================
Available trainers:
1. sft - Supervised Fine-Tuning (most common)
2. dpo - Direct Preference Optimization
3. orpo - Odds Ratio Preference Optimization
4. ppo - Proximal Policy Optimization (RL)
5. default - Generic training
Select trainer [1-5, default: 1]:
What each trainer does:
- SFT (Supervised Fine-Tuning): Train on text completion tasks, instruction following, chat
- DPO (Direct Preference Optimization): Train with preference pairs (chosen vs rejected responses)
- ORPO (Odds Ratio Preference Optimization): Alternative preference learning method
- PPO (Proximal Policy Optimization): Reinforcement learning with reward model
Provide essential project settings:
📋 Step 2: Basic Configuration
============================================================
Project name [my-llm-project]:
Specify your training data:
📋 Step 3: Dataset Configuration
============================================================
Dataset location:
• Enter local path (e.g., ./data)
• Enter HuggingFace dataset ID (e.g., timdettmers/openassistant-guanaco)
Dataset path: ./data/train.jsonl
Training split [train]: train
Validation split (optional, press Enter to skip):
Maximum samples (optional, for testing/debugging):
The wizard now includes an optional max_samples parameter that allows you to limit the number of training samples used. This is useful for:
- Quick testing with a small subset of data
- Rapid prototyping without full dataset downloads
- CI/CD pipeline testing
- Development environment validation
Simply press Enter to use the full dataset, or enter a number (e.g., 100) to limit training to that many samples.
Column Mapping (varies by trainer):
- LLM – SFT: Requires
text_columncontaining the text to train on - LLM – DPO/ORPO: Requires three columns:
prompt_text_column: The instruction/questiontext_column(chosen): The preferred responserejected_text_column: The non-preferred response
- LLM – PPO: Requires
text_columnandrl_reward_model_path - Text Classification:
text_column,target_column - Token Classification:
tokens_column,tags_column - Tabular:
target_columns, plus optionalcategorical_columns/numerical_columns - Image Classification / Regression:
image_column,target_column - Seq2Seq:
text_column(source),target_column(target), optionalmax_target_length - Extractive QA:
text_column(context),question_column,answer_column - Sentence Transformers:
sentence1_column, optionalsentence2_column/sentence3_column,target_column - VLM:
text_column,image_column, optionalprompt_text_column
Choose from popular models or enter a custom one:
📋 Step 4: Model Selection
============================================================
Popular models for sft:
1. meta-llama/Llama-3.2-1B
2. meta-llama/Llama-3.2-3B
3. Qwen/Qwen2.5-0.5B
4. Qwen/Qwen2.5-1.5B
5. google/gemma-2-2b
6. Enter custom model
Select model [1-6, default: 1]:
For custom models:
Enter model name or HF model ID: my-org/custom-model-7b
Configure advanced parameters organized by category:
📋 Step 5: Advanced Configuration (Optional)
============================================================
Would you like to configure advanced parameters?
• Training hyperparameters (learning rate, batch size, etc.)
• PEFT/LoRA settings
• Model quantization
• And more...
Configure advanced parameters? [y/N]: n
If you choose yes, you'll be prompted for parameter groups:
────────────────────────────────────────────────────────────
⚙️ Training Hyperparameters
────────────────────────────────────────────────────────────
Configure Training Hyperparameters parameters? [y/N]: y
lr
(Learning rate for training)
lr [3e-05]:
epochs
(Number of training epochs)
epochs [1]:
batch_size
(Batch size for training)
batch_size [2]:
Available parameter groups:
- Basic
- Data Processing
- Training Configuration
- Training Hyperparameters
- PEFT/LoRA
- DPO/ORPO (trainer-specific)
- Hub Integration
- Knowledge Distillation
- Hyperparameter Sweep
- Enhanced Evaluation
- Reinforcement Learning (PPO) (trainer-specific)
- Advanced Features
Review your complete configuration:
============================================================
📋 Configuration Summary
============================================================
Basic Configuration:
• project_name: my-sft-project
• trainer: sft
Dataset:
• data_path: ./data/train.jsonl
• train_split: train
• text_column: text
Model & Training:
• model: meta-llama/Llama-3.2-1B
• epochs: 1
• batch_size: 2
────────────────────────────────────────────────────────────
✓ Configuration is valid!
============================================================
🚀 Start training with this configuration? [Y/n]:
After confirmation, training begins immediately!
$ aitraining
Select trainer: 1 (SFT)
Project name: my-chat-model
Dataset path: ./data/conversations.jsonl
Training split: train
Validation split: <Enter>
Text column: text
Select model: 1 (Llama-3.2-1B)
Configure advanced? n
Start training? y
🚀 Starting training...
✓ Training job started! Job ID: 12345$ aitraining llm --interactive
Select trainer: 2 (DPO)
Project name: preference-tuned-model
Dataset path: username/preference-dataset
Training split: train
Validation split: test
Prompt column: prompt
Chosen column: chosen
Rejected column: rejected
Select model: 1 (Llama-3.2-1B)
Configure advanced? y
Configure PEFT/LoRA? y
peft: y
lora_r: 16
lora_alpha: 32
Start training? y$ aitraining
Select trainer: 4 (PPO)
Project name: rl-optimized-model
Dataset path: ./rl_prompts.jsonl
Training split: train
Text column: prompt
Reward model path: models/reward-model-7b
Select model: 1 (Llama-3.2-1B)
Configure advanced? n
Start training? yMost default values are sensible for initial experiments:
batch_size=2: Small batch size for memory efficiencyepochs=1: Quick test runlr=3e-5: Standard learning ratequantization=int4: Memory-efficient quantization
For production training, you'll likely want to configure advanced parameters.
- Local files: Use relative (
./data) or absolute paths (/home/user/data) - HuggingFace datasets: Use format
username/dataset-nameororg/dataset-name - Supported formats: JSONL, CSV, Parquet
Ensure your dataset has the required columns:
SFT:
{"text": "This is the training text..."}DPO/ORPO:
{
"prompt": "What is AI?",
"chosen": "AI stands for Artificial Intelligence...",
"rejected": "AI is a computer."
}PPO:
{"text": "Write a story about..."}When to configure advanced parameters:
- PEFT/LoRA: For memory-efficient fine-tuning of large models
- Training Hyperparameters: To tune learning rate, batch size, epochs
- Quantization: To reduce memory usage (4-bit/8-bit)
- Hub Integration: To push model to HuggingFace Hub
- Enhanced Evaluation: For detailed metrics and benchmarks
The wizard validates your configuration using Pydantic models. Common validation errors:
- Missing reward model (PPO): Must provide
rl_reward_model_path - Missing columns (DPO/ORPO): Must specify all three columns
- Invalid paths: Dataset path must exist or be valid HF dataset ID
The wizard will show these errors and let you fix them before training.
Press Ctrl+C at any time to cancel the wizard:
^C
❌ Setup cancelled.
No training will start if you cancel.
The wizard generates the same configuration as the CLI flags:
# Wizard equivalent to:
aitraining llm --train \
--trainer sft \
--project-name my-project \
--data-path ./data \
--model meta-llama/Llama-3.2-1B \
--text-column textYou can also pre-fill values via CLI and let the wizard handle the rest:
# Set trainer and model, wizard fills in the rest
aitraining llm --interactive --trainer dpo --model gpt2The wizard includes comprehensive test coverage:
pytest tests/cli/test_interactive_wizard.py -vTests verify:
- Basic wizard flow for all trainer types
- Column mapping for each trainer
- Custom model selection
- Validation and error handling
- Cancellation behavior
- Integration with
LLMTrainingParamsandAutoTrainProject
The wizard is implemented in src/autotrain/cli/interactive_wizard.py as the InteractiveWizard class:
- Reuses existing metadata:
FIELD_GROUPS,FIELD_SCOPES,get_field_info() - Validates with Pydantic: Uses
LLMTrainingParamsfor validation - Integrates seamlessly: Returns config dict compatible with
AutoTrainProject
-
Base command (
src/autotrain/cli/autotrain.py):- Detects no subcommand → launches wizard
- Creates
LLMTrainingParamsandAutoTrainProjectautomatically
-
LLM subcommand (
src/autotrain/cli/run_llm.py):--interactiveflag → launches wizard explicitly--trainwith missing params → auto-launches wizard- Merges wizard results into
self.args
To add new trainers or parameters:
- Update
FIELD_GROUPSandFIELD_SCOPESinrun_llm.py - Add trainer to
STARTER_MODELSininteractive_wizard.py - Update
_prompt_column_mapping()for trainer-specific columns - Add tests in
test_interactive_wizard.py
Wizard doesn't launch:
- Check you're running in a TTY:
sys.stdout.isatty()must beTrue - Piped commands show help instead:
aitraining | catwon't launch wizard
Validation errors:
- Read error messages carefully - they indicate which fields are missing/invalid
- Check your dataset has required columns
- Ensure reward model path exists (for PPO)
Training doesn't start:
- Make sure you confirmed with 'y' at the final prompt
- Check logs for error messages
- Verify dataset path is accessible