# /// script
# dependencies = ["trl>=0.12.0", "peft>=0.7.0", "transformers>=4.45.0", "datasets", "accelerate", "torch"]
# ///
"""Fine-tune Qwen3-0.6B on CodeForces-CoTS (100 examples)"""

import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"

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

print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")
    print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")

# Load 100 examples
print("\nLoading dataset...")
dataset = load_dataset("open-r1/codeforces-cots", "solutions", split="train").select(range(100))
print(f"Dataset: {len(dataset)} examples")

# Format for SFT: use the 'messages' column directly (already in chat format)
# The dataset has 'messages' column with list of {'role': ..., 'content': ...}
print(f"Sample messages: {dataset[0]['messages'][:1]}")

# Split: 90 train, 10 val
splits = dataset.train_test_split(test_size=0.1, seed=42)
train_ds, val_ds = splits["train"], splits["test"]
print(f"Train: {len(train_ds)}, Val: {len(val_ds)}")

peft_config = LoraConfig(
    r=8,
    lora_alpha=16,
    lora_dropout=0.05,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    bias="none",
    task_type="CAUSAL_LM"
)

# 90 examples, batch=1, accum=4 -> ~22 steps/epoch
training_args = SFTConfig(
    output_dir="./qwen3-0.6b-codeforces-cots",
    num_train_epochs=1,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,
    learning_rate=2e-4,
    warmup_ratio=0.1,
    logging_steps=2,
    logging_first_step=True,
    save_strategy="no",
    eval_strategy="steps",
    eval_steps=5,
    max_length=1024,
    push_to_hub=True,
    hub_model_id="gilbaes/qwen3-0.6b-codeforces-cots",
    report_to="none",
    bf16=True,
    gradient_checkpointing=True,
    optim="adamw_torch_fused",
    dataset_text_field=None,  # Use messages format
)

print("\nInitializing trainer...")
trainer = SFTTrainer(
    model="Qwen/Qwen3-0.6B",
    train_dataset=train_ds,
    eval_dataset=val_ds,
    peft_config=peft_config,
    args=training_args,
)

print(f"Trainable params: {trainer.model.num_parameters(only_trainable=True):,}")
print(f"Total params: {trainer.model.num_parameters():,}")

print("\n" + "="*50)
print("TRAINING START")
print("="*50 + "\n")

trainer.train()

print("\n" + "="*50)
print("PUSHING TO HUB")
print("="*50)
trainer.push_to_hub()
print("\nDone! Model: https://huggingface.co/gilbaes/qwen3-0.6b-codeforces-cots")