"""
Ult1.0 Fine-Tuning Script
==========================
Fine-tune Ult1.0 on your own data using LoRA.
Requires a GPU with ~8 GB VRAM.

Usage:
    python train.py                             # train on Alpaca
    python train.py --dataset your/dataset       # custom dataset
    python train.py --lr 1e-4 --epochs 5        # custom params
"""
import torch, argparse, os
from transformers import (
    AutoModelForCausalLM, AutoTokenizer,
    TrainingArguments, Trainer, DataCollatorForSeq2Seq
)
from peft import LoraConfig, get_peft_model, TaskType
from datasets import load_dataset

parser = argparse.ArgumentParser()
parser.add_argument("--model", default="teolm30/Ult1.0")
parser.add_argument("--dataset", default="yahma/alpaca-cleaned")
parser.add_argument("--lr", type=float, default=2e-4)
parser.add_argument("--epochs", type=int, default=3)
parser.add_argument("--batch_size", type=int, default=4)
parser.add_argument("--max_length", type=int, default=512)
parser.add_argument("--output", default="./ult10_finetuned")
args = parser.parse_args()

os.makedirs(args.output, exist_ok=True)

print(f"Loading model: {args.model}")
model = AutoModelForCausalLM.from_pretrained(
    args.model, torch_dtype=torch.bfloat16, device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(args.model)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

lora_config = LoraConfig(
    r=8, lora_alpha=16,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    task_type=TaskType.CAUSAL_LM,
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

print(f"Loading dataset: {args.dataset}")
dataset = load_dataset(args.dataset, split="train")

def format_example(ex):
    inp = f"\nInput: {ex['input']}" if ex.get("input") else ""
    return {"text": f"Instruction: {ex['instruction']}{inp}\nResponse: {ex['output']}"}

dataset = dataset.map(format_example)

def tokenize(examples):
    return tokenizer(
        examples["text"], truncation=True,
        max_length=args.max_length, padding="max_length"
    )

remove_cols = [c for c in dataset.column_names if c != "text"]
dataset = dataset.map(tokenize, remove_columns=remove_cols, batched=True)

training_args = TrainingArguments(
    output_dir=args.output,
    per_device_train_batch_size=args.batch_size,
    gradient_accumulation_steps=4,
    num_train_epochs=args.epochs,
    learning_rate=args.lr,
    logging_steps=10,
    save_strategy="epoch",
    bf16=True,
    report_to="none",
    dataloader_num_workers=4,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    data_collator=DataCollatorForSeq2Seq(tokenizer, pad_to_multiple_of=8),
)

trainer.train()
model.save_pretrained(args.output)
tokenizer.save_pretrained(args.output)
print(f"Model saved to {args.output}")