#!/usr/bin/env python3
"""Test data loading"""

import sys
from pathlib import Path
sys.path.append(str(Path(__file__).parent))

from src.training import ProgressiveTrainer
from src.progressive_model import ProgressiveReasoningModel
import yaml

# Load config
with open("config/training_config.yaml") as f:
    config = yaml.safe_load(f)

# Create dummy model wrapper
class DummyModelWrapper:
    def __init__(self):
        self.tokenizer = None
        
model_wrapper = DummyModelWrapper()

# Create trainer
trainer = ProgressiveTrainer(model_wrapper, config)

# Test data loading
stage_config = config["progressive_stages"][0]
dataset_path = stage_config["dataset_path"]
print(f"Loading dataset from: {dataset_path}")

dataset = trainer.load_dataset(dataset_path)
print(f"Loaded {len(dataset)} examples")

if len(dataset) > 0:
    print(f"First example: {dataset[0]}")