"""Ult1-Coding GPU fine-tuning script.
Usage: python train_coding.py (requires 8+ GB VRAM GPU)
"""
import json, torch, warnings
from datasets import Dataset
from transformers import (
    AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
)
from peft import LoraConfig, get_peft_model, TaskType
from huggingface_hub import HfApi
warnings.filterwarnings("ignore")

import os
HF_TOKEN = os.getenv("HF_TOKEN", "")
MODEL_ID = "Qwen/Qwen2.5-3B-Instruct"
REPO_ID = "teolm30/Ult1-coding"

def format_example(example):
    return {
        "text": f"<|im_start|>user\n{example['instruction']}<|im_end|>\n<|im_start|>assistant\n{example['response']}<|im_end|>\n"
    }

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
tokenizer.pad_token = tokenizer.eos_token

print("Loading training data from HF...")
api = HfApi()
data_path = api.hf_hub_download(repo_id=REPO_ID, filename="training_data.json", token=HF_TOKEN)
with open(data_path, "r") as f:
    raw_data = json.load(f)

formatted = [format_example(ex) for ex in raw_data]
dataset = Dataset.from_list(formatted)

def tokenize_fn(examples):
    result = tokenizer(
        examples["text"], truncation=True, max_length=512,
        padding="max_length", return_tensors=None
    )
    result["labels"] = result["input_ids"].copy()
    return result

dataset = dataset.map(tokenize_fn, remove_columns=["text"], batched=True)
dataset = dataset.train_test_split(test_size=0.1)

print(f"Training samples: {len(dataset['train'])}, Validation: {len(dataset['test'])}")

print("Loading model...")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, torch_dtype=torch.bfloat16,
    device_map="auto", attn_implementation="flash_attention_2",
)

model.train()
lora_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"],
    bias="none", task_type=TaskType.CAUSAL_LM,
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

training_args = TrainingArguments(
    output_dir="./ult1_coding_trained",
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,
    learning_rate=3e-4,
    num_train_epochs=3,
    logging_steps=1,
    save_strategy="epoch",
    evaluation_strategy="epoch",
    bf16=True,
    gradient_checkpointing=True,
    logging_dir="./logs",
)

trainer = Trainer(
    model=model, args=training_args,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
)

trainer.train()
model.save_pretrained("./ult1_coding_trained/final", safe_serialization=True)
tokenizer.save_pretrained("./ult1_coding_trained/final")

api.upload_folder(
    folder_path="./ult1_coding_trained/final",
    repo_id=REPO_ID, token=HF_TOKEN,
    commit_message="Ult1-Coding trained adapter - GPU fine-tuned on coding data",
    ignore_patterns=["*.pt", "checkpoint-*"],
)
print(f"Trained model uploaded to https://huggingface.co/{REPO_ID}")