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

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

# Load dataset with editorials for better instruction following
dataset = load_dataset(
    "open-r1/codeforces-cots",
    name="solutions_w_editorials_decontaminated",
    split="train"
)

# Create train/eval split (90/10)
dataset_split = dataset.train_test_split(test_size=0.1, seed=42)

# LoRA configuration for efficient fine-tuning
peft_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    task_type="CAUSAL_LM"
)

# SFT Training configuration
training_args = SFTConfig(
    output_dir="qwen3-0.6b-codeforces-instruct",
    
    # Training hyperparameters
    num_train_epochs=3,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    gradient_accumulation_steps=4,  # Effective batch size: 16
    gradient_checkpointing=True,
    
    # Learning rate and optimization
    learning_rate=2e-4,
    lr_scheduler_type="cosine",
    warmup_ratio=0.1,
    optim="paged_adamw_8bit",
    
    # Evaluation and logging
    eval_strategy="steps",
    eval_steps=100,
    logging_steps=10,
    save_strategy="steps",
    save_steps=200,
    save_total_limit=3,
    
    # Hub integration (CRITICAL - saves model to Hub)
    push_to_hub=True,
    hub_model_id="kneeraj/qwen3-0.6b-codeforces-instruct",
    hub_strategy="every_save",
    hub_private_repo=False,
    
    # Trackio monitoring
    report_to="trackio",
    project="codeforces-finetuning",
    run_name="qwen3-0.6b-codeforces-sft",
    
    # Performance optimizations
    bf16=True,
    max_grad_norm=1.0,
    
    # Data processing
    max_seq_length=2048,  # CodeForces problems can be lengthy
    dataset_text_field="messages",  # Use chat format
    packing=False,  # Don't pack for instruction following
)

# Initialize trainer
trainer = SFTTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",  # Using Qwen2.5-0.5B as base (Qwen3-0.6B may not be available)
    train_dataset=dataset_split["train"],
    eval_dataset=dataset_split["test"],
    peft_config=peft_config,
    args=training_args,
)

print("Starting training...")
print(f"Training samples: {len(dataset_split['train'])}")
print(f"Evaluation samples: {len(dataset_split['test'])}")

# Train the model
trainer.train()

# Final push to Hub
print("Pushing final model to Hub...")
trainer.push_to_hub()

print("Training complete! Model saved to: kneeraj/qwen3-0.6b-codeforces-instruct")