# /// script
# dependencies = ["trl>=0.12.0", "peft>=0.7.0", "trackio", "torch", "transformers"]
# ///

from datasets import load_dataset
from peft import LoraConfig
from trl import SFTTrainer, SFTConfig
import trackio

print("🚀 Starting quick proof-of-concept training...")

# Load tiny subset for quick test
dataset = load_dataset("trl-lib/Capybara", split="train[:50]")

print(f"📊 Dataset loaded: {len(dataset)} examples")

# LoRA configuration
peft_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
    task_type="CAUSAL_LM"
)

# Training configuration
training_args = SFTConfig(
    output_dir="comfyui-specialist-test",
    num_train_epochs=1,
    max_steps=50,  # Just 50 steps for quick validation
    per_device_train_batch_size=2,
    gradient_accumulation_steps=4,
    learning_rate=2e-4,
    logging_steps=5,
    save_strategy="steps",
    save_steps=25,
    push_to_hub=True,
    hub_model_id="lokegud/comfyui-specialist-test",
    hub_strategy="every_save",
    report_to="trackio",
    project="comfyui-specialist",
    run_name="quick-test",
    gradient_checkpointing=True,
)

print("🔧 Initializing trainer...")

# Initialize trainer
trainer = SFTTrainer(
    model="Qwen/Qwen2.5-0.5B",
    train_dataset=dataset,
    peft_config=peft_config,
    args=training_args,
)

print("🏋️ Training...")
trainer.train()

print("📤 Pushing to Hub...")
trainer.push_to_hub()

print("✅ Quick test complete! Model saved to: lokegud/comfyui-specialist-test")