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

from datasets import load_dataset, concatenate_datasets
from peft import LoraConfig, PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import SFTTrainer, SFTConfig
import trackio
import torch

print("Loading base model and merging SFT adapter...")

# Load the base Qwen model
base_model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True
)

# Load and merge the existing SFT adapter to preserve code capabilities
model = PeftModel.from_pretrained(base_model, "chaddy81/qwen3-0.6b-multicode-sft")
model = model.merge_and_unload()  # Merge adapter into base model

print("SFT adapter merged successfully!")

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B", trust_remote_code=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

print("Loading datasets...")

# Load the validated Tailwind dataset (already in messages format)
tailwind_ds = load_dataset("summykai/tailwind-v4-sft-mix-001", split="train")
print(f"Tailwind dataset: {len(tailwind_ds)} examples")

# Load the high-quality UI generation dataset
uigen_ds = load_dataset("smirki/UIGEN-T1.1-TAILWIND", split="train")
print(f"UIGEN dataset: {len(uigen_ds)} examples")

# Format UIGEN dataset to messages format
def format_uigen_to_messages(example):
    # Extract code from markdown code blocks if present
    answer = example["answer"]
    if answer.startswith("```html"):
        answer = answer[7:]  # Remove ```html
    if answer.endswith("```"):
        answer = answer[:-3]  # Remove trailing ```

    # Include the reasoning as part of a richer response
    response = answer.strip()

    return {
        "messages": [
            {"role": "user", "content": example["question"]},
            {"role": "assistant", "content": response}
        ]
    }

print("Formatting UIGEN dataset...")
uigen_formatted = uigen_ds.map(
    format_uigen_to_messages,
    remove_columns=uigen_ds.column_names
)

# Combine datasets
print("Combining datasets...")
combined_dataset = concatenate_datasets([tailwind_ds, uigen_formatted])
combined_dataset = combined_dataset.shuffle(seed=42)

print(f"Total training examples: {len(combined_dataset)}")

# Create train/eval split
dataset_split = combined_dataset.train_test_split(test_size=0.05, seed=42)
print(f"Train: {len(dataset_split['train'])}, Eval: {len(dataset_split['test'])}")

print("Initializing trainer...")
trainer = SFTTrainer(
    model=model,  # Use the merged model
    processing_class=tokenizer,
    train_dataset=dataset_split["train"],
    eval_dataset=dataset_split["test"],
    peft_config=LoraConfig(
        r=16,  # Reduced from 32 to save memory
        lora_alpha=32,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        lora_dropout=0.05,
        bias="none",
        task_type="CAUSAL_LM"
    ),
    args=SFTConfig(
        output_dir="qwen3-0.6b-design-sft",
        push_to_hub=True,
        hub_model_id="chaddy81/qwen3-0.6b-design-sft",
        num_train_epochs=3,
        per_device_train_batch_size=1,  # Reduced from 2
        gradient_accumulation_steps=16,  # Increased to maintain effective batch size
        learning_rate=2e-4,
        lr_scheduler_type="cosine",
        warmup_ratio=0.1,
        max_length=2048,  # Reduced from 4096 to save memory
        logging_steps=10,
        eval_strategy="no",  # Skip eval during training to save memory
        save_strategy="steps",
        save_steps=300,
        bf16=True,
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},  # More memory efficient
        optim="adamw_8bit",  # 8-bit optimizer to save memory
        report_to="trackio",
        run_name="qwen3-design-sft-v2",
    )
)

print("Starting training...")
trainer.train()

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

print("Training complete!")