fix launch script

.cursorrules

@@ -1,277 +0,0 @@
----
-description: SmolLM3 Fine-tuning Pipeline - Project Rules and Conventions
-globs: ["**/*.py", "**/*.sh", "**/*.md", "**/*.json"]
-alwaysApply: true
----
-
-# SmolLM3 Fine-tuning Pipeline Project Rules
-
-## Project Overview
-This is a comprehensive end-to-end fine-tuning pipeline for SmolLM3 models with Trackio monitoring, Hugging Face integration, and interactive configuration management.
-
-## Core Architecture
-
-### Directory Structure
-- `config/` - Training configuration files for different scenarios
-- `src/` - Core training and model logic
-- `scripts/` - Utility scripts for deployment, dataset management, and model pushing
-- `docs/` - Comprehensive documentation and guides
-- `templates/` - Templates for HF Spaces and datasets
-- `tests/` - Test files and debugging scripts
-- `outputs/` - Training outputs and checkpoints
-
-### Key Components
-
-#### Training Configurations
-- **Basic Training**: SmolLM3-3B + OpenHermes-FR, 3 epochs, batch size 2
-- **H100 Lightweight**: SmolLM3-3B + OpenHermes-FR (80K samples), 1 epoch, batch size 16
-- **A100 Large Scale**: SmolLM3-3B + OpenHermes-FR, 1.3 passes, batch size 8
-- **Multiple Passes**: SmolLM3-3B + OpenHermes-FR, 4 epochs, batch size 6
-- **Custom Configuration**: User-defined parameters
-
-#### Core Modules
-- `src/train.py` - Main training orchestration
-- `src/model.py` - Model loading and configuration
-- `src/data.py` - Dataset processing and loading
-- `src/monitoring.py` - Trackio integration and metrics
-- `src/trainer.py` - Training loop and optimization
-
-## Coding Conventions
-
-### Python Style
-- Use type hints for all function parameters and return values
-- Follow PEP 8 for formatting
-- Use descriptive variable names in snake_case
-- Add comprehensive docstrings for all functions
-- Use f-strings for string formatting
-
-### Configuration Management
-- All training configs inherit from `SmolLM3Config` base class
-- Use dataclasses for configuration objects
-- Validate configuration parameters in __post_init__
-- Support both YAML and Python configuration files
-
-### Error Handling
-- Use try-except blocks for external API calls (HF, Trackio)
-- Log errors with appropriate context
-- Provide user-friendly error messages
-- Implement graceful degradation for optional features
-
-### Monitoring Integration
-- Always include Trackio URL and experiment name in configs
-- Log metrics every N steps (configurable)
-- Save checkpoints and artifacts to HF Datasets
-- Use structured logging with consistent field names
-
-## File Naming Conventions
-
-### Configuration Files
-- `train_smollm3_*.py` - Training configurations
-- `*_config.py` - General configuration files
-- Use descriptive suffixes: `_h100_lightweight`, `_a100_large`, `_multiple_passes`
-
-### Script Files
-- `deploy_*.py` - Deployment scripts
-- `setup_*.py` - Setup and initialization scripts
-- `push_*.py` - Model pushing scripts
-- `configure_*.py` - Configuration scripts
-
-### Test Files
-- `test_*.py` - Test files
-- `debug_*.py` - Debugging scripts
-- Include descriptive names indicating what they test
-
-## Training Pipeline Workflow
-
-### Interactive Pipeline (`launch.sh`)
-1. **Authentication**: HF username and token validation
-2. **Configuration Selection**: Choose from predefined configs or custom
-3. **Experiment Setup**: Configure experiment name and repositories
-4. **Environment Setup**: Install dependencies and setup virtual environment
-5. **Deployment**: Deploy Trackio Space and setup HF Dataset
-6. **Training**: Execute training with monitoring
-7. **Model Push**: Upload model to HF Hub with documentation
-8. **Testing**: Validate uploaded model functionality
-
-### Configuration Selection Logic
-- Basic Training: Default for beginners and learning
-- H100 Lightweight: Rapid experiments on H100 GPUs
-- A100 Large Scale: Serious research and production
-- Multiple Passes: Thorough training for production models
-- Custom: User-defined parameters for specific needs
-
-## Dataset Management
-
-### Supported Formats
-- Hugging Face Datasets format
-- JSON files with prompt/completion pairs
-- Chat format with messages array
-- Custom formats with conversion functions
-
-### Dataset Processing
-- Automatic format detection and conversion
-- Random sampling for lightweight configurations
-- Validation split creation
-- Bad entry filtering and handling
-
-### Dataset Sampling (H100 Lightweight)
-- 80,000 random samples from OpenHermes-FR
-- 1,000 validation samples (if available)
-- Fixed random seed (42) for reproducibility
-- Automatic sampling during dataset preparation
-
-## Model Management
-
-### Model Loading
-- Support for HuggingFaceTB/SmolLM3-3B
-- Flash attention and gradient checkpointing
-- Mixed precision training (fp16/bf16)
-- Device mapping and memory optimization
-
-### Model Pushing
-- Comprehensive model cards with training details
-- Automatic README generation
-- License and usage information
-- Training metrics and configuration
-
-## Monitoring and Tracking
-
-### Trackio Integration
-- Real-time metrics logging
-- Training curves visualization
-- Resource usage monitoring
-- Artifact storage and versioning
-
-### Metrics to Track
-- Training and validation loss
-- Learning rate schedule
-- Gradient norms
-- GPU utilization and memory
-- Training speed (steps/second)
-
-## Error Handling and Validation
-
-### Input Validation
-- Validate HF tokens before use
-- Check CUDA availability
-- Verify dataset accessibility
-- Validate configuration parameters
-
-### Error Recovery
-- Graceful handling of network issues
-- Automatic retry for failed operations
-- Checkpoint recovery for interrupted training
-- Fallback options for optional features
-
-## Documentation Standards
-
-### README Files
-- Clear project description
-- Installation instructions
-- Usage examples
-- Configuration options
-- Troubleshooting guide
-
-### Code Documentation
-- Comprehensive docstrings
-- Type hints for all functions
-- Example usage in docstrings
-- Parameter descriptions
-- Return value documentation
-
-## Testing and Validation
-
-### Test Categories
-- Unit tests for core functions
-- Integration tests for pipeline
-- Configuration validation tests
-- Model loading and saving tests
-- Dataset processing tests
-
-### Debugging Tools
-- Standalone test scripts
-- Configuration validation
-- Model testing utilities
-- Dataset inspection tools
-
-## Performance Optimization
-
-### H100 Optimizations
-- Larger batch sizes (16 vs 8 for A100)
-- Reduced gradient accumulation (4 vs 16)
-- Higher learning rates (8e-6 vs 5e-6)
-- Optimized data loading (4 workers, pin memory)
-
-### Memory Management
-- Gradient checkpointing for large models
-- Mixed precision training
-- Dynamic batch sizing
-- Memory-efficient data loading
-
-## Security and Best Practices
-
-### Token Management
-- Never hardcode tokens in code
-- Use environment variables
-- Validate tokens before use
-- Secure token storage
-
-### Data Privacy
-- Filter sensitive data from datasets
-- Validate dataset contents
-- Secure data transmission
-- Proper data disposal
-
-## Deployment and CI/CD
-
-### Environment Setup
-- Python virtual environments
-- CUDA-compatible PyTorch
-- Required dependencies installation
-- System package management
-
-### Automated Deployment
-- Trackio Space deployment
-- HF Dataset setup
-- Model repository creation
-- Configuration file generation
-
-## Troubleshooting Guidelines
-
-### Common Issues
-- CUDA out of memory: Reduce batch size
-- Network timeouts: Check internet connection
-- Token validation: Verify HF token permissions
-- Dataset loading: Check dataset accessibility
-
-### Debugging Steps
-1. Check system requirements
-2. Validate configuration
-3. Test individual components
-4. Review logs and error messages
-5. Verify external service connectivity
-
-## Future Enhancements
-
-### Planned Features
-- Multi-GPU training support
-- Advanced dataset sampling strategies
-- Automated hyperparameter optimization
-- Enhanced monitoring and visualization
-- Support for additional model architectures
-
-### Extensibility
-- Modular configuration system
-- Plugin architecture for custom features
-- Support for custom datasets and models
-- Flexible monitoring integration
-
----
-
-**When working with this codebase:**
-- Always consider the end-to-end pipeline workflow
-- Follow the established configuration patterns
-- Include proper error handling and validation
-- Maintain comprehensive documentation
-- Test changes thoroughly before deployment
-- Consider performance implications for different hardware configurations 

launch.sh

@@ -489,113 +489,45 @@ echo "==========================================="
 cd ../..
 create_training_config "$CONFIG_FILE"
 
-# Step 13: Download and prepare dataset
-print_step "Step 13: Preparing Dataset"
-echo "==============================="
+# Step 13: Dataset preparation (handled by src/data.py during training)
+print_step "Step 13: Dataset Configuration"
+echo "=================================="
 
-python -c "
-from datasets import load_dataset
-import json
-import os
-import random
-
-# Load dataset
-print('Loading dataset: $DATASET_NAME')
-dataset = load_dataset('$DATASET_NAME')
-
-# Create dataset directory
-os.makedirs('training_dataset', exist_ok=True)
-
-# Convert to training format
-def convert_to_training_format(example):
-    # Handle different dataset formats
-    if 'prompt' in example and 'completion' in example:
-        return {
-            'prompt': example['prompt'],
-            'completion': example['completion']
-        }
-    elif 'instruction' in example and 'output' in example:
-        return {
-            'prompt': example['instruction'],
-            'completion': example['output']
-        }
-    elif 'messages' in example:
-        # Handle chat format
-        messages = example['messages']
-        if len(messages) >= 2:
-            return {
-                'prompt': messages[0]['content'],
-                'completion': messages[1]['content']
-            }
-    else:
-        # Fallback
-        return {
-            'prompt': str(example.get('input', '')),
-            'completion': str(example.get('output', ''))
-        }
-
-# Process train split
-train_data = []
-for example in dataset['train']:
-    training_example = convert_to_training_format(example)
-    if training_example['prompt'] and training_example['completion']:
-        train_data.append(training_example)
-
-# Apply dataset sampling for lightweight configuration
-if '$TRAINING_CONFIG_TYPE' == 'H100 Lightweight (Rapid)' and len(train_data) > ${DATASET_SAMPLE_SIZE:-0}:
-    print(f'Sampling {${DATASET_SAMPLE_SIZE:-80000}} random samples from {len(train_data)} total samples')
-    random.seed(42)  # For reproducibility
-    train_data = random.sample(train_data, ${DATASET_SAMPLE_SIZE:-80000})
-    print(f'Selected {len(train_data)} samples for lightweight training')
-
-# Process validation split if available
-val_data = []
-if 'validation' in dataset:
-    for example in dataset['validation']:
-        training_example = convert_to_training_format(example)
-        if training_example['prompt'] and training_example['completion']:
-            val_data.append(training_example)
-
-# For lightweight config, also sample validation if it's large
-if '$TRAINING_CONFIG_TYPE' == 'H100 Lightweight (Rapid)' and len(val_data) > 1000:
-    print(f'Sampling 1000 random validation samples from {len(val_data)} total')
-    random.seed(42)  # For reproducibility
-    val_data = random.sample(val_data, 1000)
-
-# Save to files
-with open('training_dataset/train.json', 'w') as f:
-    json.dump(train_data, f, indent=2)
-
-if val_data:
-    with open('training_dataset/validation.json', 'w') as f:
-        json.dump(val_data, f, indent=2)
-
-print(f'Dataset prepared: {len(train_data)} train samples, {len(val_data)} validation samples')
-"
+print_info "Dataset will be loaded directly by src/data.py during training"
+print_info "Dataset: $DATASET_NAME"
+if [ "$TRAINING_CONFIG_TYPE" = "H100 Lightweight (Rapid)" ]; then
+    print_info "Sample size: ${DATASET_SAMPLE_SIZE:-80000} (will be handled by data.py)"
+fi
 
 # Step 14: Calculate training parameters
 print_step "Step 14: Calculating Training Parameters"
 echo "============================================"
 
-TOTAL_SAMPLES=$(python -c "import json; data=json.load(open('training_dataset/train.json')); print(len(data))")
+# Estimate training steps
 EFFECTIVE_BATCH_SIZE=$((BATCH_SIZE * GRADIENT_ACCUMULATION_STEPS))
-STEPS_PER_EPOCH=$((TOTAL_SAMPLES / EFFECTIVE_BATCH_SIZE))
-MAX_STEPS=$((STEPS_PER_EPOCH * MAX_EPOCHS))
-
-echo "  Total samples: $TOTAL_SAMPLES"
 echo "  Effective batch size: $EFFECTIVE_BATCH_SIZE"
-echo "  Steps per epoch: $STEPS_PER_EPOCH"
-echo "  Total training steps: $MAX_STEPS"
+echo "  Learning rate: $LEARNING_RATE"
+echo "  Max epochs: $MAX_EPOCHS"
+echo "  Sequence length: $MAX_SEQ_LENGTH"
+echo "  Training steps will be calculated by the training script"
 
 # Step 15: Start training
 print_step "Step 15: Starting Training"
 echo "=============================="
 
-python src/train.py "$CONFIG_FILE" \
-    --dataset_dir training_dataset \
+print_info "Using existing scripts/training/train.py script with the following parameters:"
+echo "  Model: $MODEL_NAME"
+echo "  Dataset: $DATASET_NAME"
+echo "  Output: /output-checkpoint"
+echo "  Batch size: $BATCH_SIZE"
+echo "  Learning rate: $LEARNING_RATE"
+echo "  Sequence length: $MAX_SEQ_LENGTH"
+
+# Run the existing training script
+python scripts/training/train.py "$CONFIG_FILE" \
+    --dataset_dir "$DATASET_NAME" \
     --out_dir /output-checkpoint \
     --init_from scratch \
-    --max_iters $MAX_STEPS \
     --batch_size $BATCH_SIZE \
     --learning_rate $LEARNING_RATE \
     --gradient_accumulation_steps $GRADIENT_ACCUMULATION_STEPS \
@@ -613,38 +545,23 @@ python src/train.py "$CONFIG_FILE" \
 print_step "Step 16: Pushing Model to HF Hub"
 echo "====================================="
 
+print_info "Using scripts/model_tonic/push_to_huggingface.py script"
+echo "  Checkpoint: /output-checkpoint"
+echo "  Repository: $REPO_NAME"
+
+# Run the existing push script
 python scripts/model_tonic/push_to_huggingface.py /output-checkpoint "$REPO_NAME" \
     --token "$HF_TOKEN" \
     --trackio-url "$TRACKIO_URL" \
     --experiment-name "$EXPERIMENT_NAME" \
     --dataset-repo "$TRACKIO_DATASET_REPO"
 
-# Step 17: Test the uploaded model
-print_step "Step 17: Testing Uploaded Model"
-echo "==================================="
-
-python -c "
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
-
-print('Loading uploaded model...')
-model = AutoModelForCausalLM.from_pretrained('$REPO_NAME', torch_dtype=torch.float16, device_map='auto')
-tokenizer = AutoTokenizer.from_pretrained('$REPO_NAME')
-
-print('Testing model generation...')
-prompt = 'Hello, how are you?'
-inputs = tokenizer(prompt, return_tensors='pt').to(model.device)
-outputs = model.generate(**inputs, max_new_tokens=50, do_sample=True, temperature=0.7)
-response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-print(f'Prompt: {prompt}')
-print(f'Response: {response}')
-print('✅ Model test completed successfully!')
-"
-
-# Step 18: Create summary report
-print_step "Step 18: Creating Summary Report"
+# Step 17: Create summary report
+print_step "Step 17: Creating Summary Report"
 echo "===================================="
 
+
+
 cat > training_summary.md << EOF
 # SmolLM3 Fine-tuning Summary
 
@@ -665,8 +582,6 @@ fi)
 - **Gradient Accumulation**: $GRADIENT_ACCUMULATION_STEPS
 - **Learning Rate**: $LEARNING_RATE
 - **Max Epochs**: $MAX_EPOCHS
-- **Max Steps**: $MAX_STEPS
-- **Total Samples**: $TOTAL_SAMPLES
 - **Sequence Length**: $MAX_SEQ_LENGTH
 
 ## Results
@@ -682,7 +597,6 @@ fi)
 
 ## Files Created
 - Training configuration: \`$CONFIG_FILE\`
-- Dataset: \`training_dataset/\`
 - Model checkpoint: \`/output-checkpoint/\`
 - Training logs: \`training.log\`
 - Summary report: \`training_summary.md\`
@@ -690,6 +604,10 @@ EOF
 
 print_status "Summary report saved to: training_summary.md"
 
+# Clean up temporary files
+print_info "Cleaning up temporary files..."
+rm -f deploy_input.txt
+
 # Final summary
 echo ""
 print_header "🎉 End-to-End Pipeline Completed Successfully!"

scripts/trackio_tonic/deploy_trackio_space.py

@@ -30,13 +30,21 @@ class TrackioSpaceDeployer:
             cmd = [
                 "huggingface-cli", "repo", "create",
                 f"{self.username}/{self.space_name}",
-                "--type", "space",
-                "--space-sdk", "gradio",
-                "--space-hardware", "cpu-basic"
+                "--type", "space"
             ]
             
+            # Try to create the space first
             result = subprocess.run(cmd, capture_output=True, text=True)
             
+            if result.returncode != 0:
+                # Try alternative approach without space-specific flags
+                print("Retrying with basic space creation...")
+                cmd = [
+                    "huggingface-cli", "repo", "create",
+                    f"{self.username}/{self.space_name}"
+                ]
+                result = subprocess.run(cmd, capture_output=True, text=True)
+            
             if result.returncode == 0:
                 print(f"✅ Space created successfully: {self.space_url}")
                 return True

src/config.py

@@ -3,9 +3,27 @@ Configuration management for SmolLM3 fine-tuning
 """
 
 import os
+import sys
 import importlib.util
 from typing import Any
-from config.train_smollm3 import SmolLM3Config, get_config as get_default_config
+
+# Add the project root to Python path
+project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+if project_root not in sys.path:
+    sys.path.insert(0, project_root)
+
+# Add config directory to path
+config_dir = os.path.join(project_root, 'config')
+if config_dir not in sys.path:
+    sys.path.insert(0, config_dir)
+
+try:
+    from config.train_smollm3 import SmolLM3Config, get_config as get_default_config
+except ImportError:
+    # Fallback: try direct import
+    import sys
+    sys.path.insert(0, os.path.join(project_root, 'config'))
+    from train_smollm3 import SmolLM3Config, get_config as get_default_config
 
 def get_config(config_path: str) -> SmolLM3Config:
     """Load configuration from file or return default"""

src/train.py

@@ -16,7 +16,17 @@ from typing import Optional, Dict, Any
 # Add the current directory to the path for imports
 sys.path.append(os.path.dirname(os.path.abspath(__file__)))
 
-from config import get_config
+# Add project root to path for config imports
+project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+if project_root not in sys.path:
+    sys.path.insert(0, project_root)
+
+try:
+    from config import get_config
+except ImportError:
+    # Fallback: try direct import
+    sys.path.insert(0, os.path.join(project_root, 'src'))
+    from config import get_config
 from model import SmolLM3Model
 from data import SmolLM3Dataset
 from trainer import SmolLM3Trainer

tests/test_dataset.py

@@ -0,0 +1,88 @@
+#!/usr/bin/env python3
+"""
+Test script to verify OpenHermes-FR dataset loading
+"""
+
+from datasets import load_dataset
+import json
+import random
+
+def test_openhermes_fr():
+    """Test loading and processing OpenHermes-FR dataset"""
+    
+    print("Loading OpenHermes-FR dataset...")
+    try:
+        dataset = load_dataset('legmlai/openhermes-fr')
+        print(f"✅ Dataset loaded successfully")
+        print(f"   Train samples: {len(dataset['train'])}")
+        if 'validation' in dataset:
+            print(f"   Validation samples: {len(dataset['validation'])}")
+        
+        # Show sample structure
+        sample = dataset['train'][0]
+        print(f"\n📋 Sample structure:")
+        for key, value in sample.items():
+            if isinstance(value, str) and len(value) > 100:
+                print(f"   {key}: {value[:100]}...")
+            else:
+                print(f"   {key}: {value}")
+        
+        # Test conversion
+        print(f"\n🔄 Testing conversion...")
+        
+        def convert_to_training_format(example):
+            # Handle OpenHermes-FR format specifically
+            if 'prompt' in example and 'accepted_completion' in example:
+                return {
+                    'prompt': example['prompt'],
+                    'completion': example['accepted_completion']
+                }
+            elif 'prompt' in example and 'completion' in example:
+                return {
+                    'prompt': example['prompt'],
+                    'completion': example['completion']
+                }
+            else:
+                return None
+        
+        # Process first 10 examples
+        train_data = []
+        for i, example in enumerate(dataset['train'][:10]):
+            training_example = convert_to_training_format(example)
+            if training_example and training_example['prompt'] and training_example['completion']:
+                # Filter out bad entries
+                if 'bad_entry' in example and example['bad_entry']:
+                    print(f"   Skipping bad entry {i}")
+                    continue
+                train_data.append(training_example)
+                print(f"   ✅ Converted example {i}")
+        
+        print(f"\n📊 Conversion results:")
+        print(f"   Converted: {len(train_data)} valid examples")
+        
+        if train_data:
+            print(f"\n📝 Sample converted example:")
+            sample = train_data[0]
+            print(f"   Prompt: {sample['prompt'][:100]}...")
+            print(f"   Completion: {sample['completion'][:100]}...")
+        
+        # Test sampling
+        if len(dataset['train']) > 100:
+            print(f"\n🎲 Testing sampling...")
+            random.seed(42)
+            sampled_indices = random.sample(range(len(dataset['train'])), 5)
+            print(f"   Sampled indices: {sampled_indices}")
+        
+        return True
+        
+    except Exception as e:
+        print(f"❌ Error loading dataset: {e}")
+        return False
+
+if __name__ == "__main__":
+    success = test_openhermes_fr()
+    if success:
+        print("\n✅ Dataset test completed successfully!")
+    else:
+        print("\n❌ Dataset test failed!")
+        exit(1) 

tests/test_dataset_loading.py

@@ -0,0 +1,71 @@
+#!/usr/bin/env python3
+"""
+Test script to verify dataset loading works correctly
+"""
+
+import os
+import sys
+import json
+from datasets import load_dataset
+
+def test_dataset_loading():
+    """Test loading the OpenHermes-FR dataset"""
+    print("Testing dataset loading...")
+    
+    try:
+        # Load the dataset
+        dataset = load_dataset("legmlai/openhermes-fr")
+        print(f"✅ Dataset loaded successfully")
+        print(f"  Train samples: {len(dataset['train'])}")
+        
+        # Check the first few examples
+        print("\nSample examples:")
+        for i in range(min(3, len(dataset['train'])):
+            example = dataset['train'][i]
+            print(f"\nExample {i+1}:")
+            print(f"  Keys: {list(example.keys())}")
+            print(f"  Prompt: {example.get('prompt', 'N/A')[:100]}...")
+            print(f"  Accepted completion: {example.get('accepted_completion', 'N/A')[:100]}...")
+            print(f"  Bad entry: {example.get('bad_entry', 'N/A')}")
+        
+        # Test filtering bad entries
+        print(f"\nFiltering bad entries...")
+        original_size = len(dataset['train'])
+        filtered_dataset = dataset['train'].filter(lambda x: not x.get('bad_entry', False))
+        filtered_size = len(filtered_dataset)
+        print(f"  Original size: {original_size}")
+        print(f"  Filtered size: {filtered_size}")
+        print(f"  Removed: {original_size - filtered_size} bad entries")
+        
+        # Test conversion to training format
+        print(f"\nTesting conversion to training format...")
+        train_data = []
+        for i, example in enumerate(filtered_dataset):
+            if i >= 5:  # Just test first 5 examples
+                break
+            
+            if 'prompt' in example and 'accepted_completion' in example:
+                train_data.append({
+                    'prompt': example['prompt'],
+                    'completion': example['accepted_completion']
+                })
+        
+        print(f"  Converted {len(train_data)} examples to training format")
+        
+        # Save a small sample
+        os.makedirs('test_dataset', exist_ok=True)
+        with open('test_dataset/train.json', 'w') as f:
+            json.dump(train_data, f, indent=2)
+        
+        print(f"✅ Test completed successfully!")
+        print(f"  Sample saved to: test_dataset/train.json")
+        
+        return True
+        
+    except Exception as e:
+        print(f"❌ Dataset loading failed: {e}")
+        return False
+
+if __name__ == "__main__":
+    success = test_dataset_loading()
+    sys.exit(0 if success else 1)