"""QLoRA training script for ORTOS bot.
Uses transformers + bitsandbytes + peft + trl (no unsloth/vllm).
"""
import json
import torch
from datasets import Dataset
from transformers import (
    AutoModelForCausalLM, AutoTokenizer,
)
from trl import SFTTrainer, SFTConfig
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

MODEL_NAME = "unsloth/Meta-Llama-3.1-8B-bnb-4bit"
OUTPUT_DIR = "lora_ortos"


def format_chat(example):
    return {
        "text": f"<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n{example['instruction']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n{example['response']}<|eot_id|>"
    }


# Load dataset
with open("lora_dataset.jsonl", encoding="utf-8") as f:
    data = [json.loads(line) for line in f]

dataset = Dataset.from_list(data)
dataset = dataset.map(format_chat)
dataset = dataset.train_test_split(test_size=0.05)
print(f"Train: {len(dataset['train'])}, Eval: {len(dataset['test'])}")

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    device_map="auto",
    dtype=torch.bfloat16,
)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
tokenizer.pad_token = tokenizer.eos_token

# Prepare for k-bit training (gradient checkpointing + freezing base)
model = prepare_model_for_kbit_training(model)

# LoRA adapters
lora_config = LoraConfig(
    r=16,
    lora_alpha=16,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
    lora_dropout=0,
    bias="none",
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

# Train
trainer = SFTTrainer(
    model=model,
    processing_class=tokenizer,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    args=SFTConfig(
        output_dir=OUTPUT_DIR,
        dataset_text_field="text",
        max_length=2048,
        per_device_train_batch_size=2,
        gradient_accumulation_steps=4,
        warmup_steps=5,
        num_train_epochs=2,
        learning_rate=2e-4,
        fp16=False,
        bf16=torch.cuda.is_bf16_supported(),
        logging_steps=10,
        eval_steps=50,
        save_steps=100,
        report_to="none",
        save_total_limit=2,
    ),
)

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