# Entrenamiento SFT + LoRA sobre Qwen2.5-7B-Instruct
import torch
from datasets import load_dataset
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    BitsAndBytesConfig,
    TrainingArguments,
)
from trl import SFTTrainer
from peft import (
    LoraConfig,
    get_peft_model,
    prepare_model_for_kbit_training,
    PeftModel,
)

base_model = "Qwen/Qwen2.5-7B-Instruct"
dataset_path = "/content/drive/Othercomputers/Mi PC/Mis proyectos en code/ASISTENTE-VIRTUAL/dataset_malo_convertido_a_bueno_shuffled_clean.jsonl"

# 1) Dataset y split 90/10
ds_raw = load_dataset("json", data_files=dataset_path, split="train")
ds_split = ds_raw.train_test_split(test_size=0.1, seed=42)
train_ds, eval_ds = ds_split["train"], ds_split["test"]

# 2) Tokenizer y modelo en 4-bit
tok = AutoTokenizer.from_pretrained(base_model)
if tok.pad_token is None:
    tok.pad_token = tok.eos_token

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)

mdl = AutoModelForCausalLM.from_pretrained(
    base_model,
    quantization_config=bnb_config,
    device_map="auto",
)
mdl = prepare_model_for_kbit_training(mdl)

# 3) LoRA ligera
lora_cfg = LoraConfig(
    r=32,
    lora_alpha=64,
    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",
)
mdl = get_peft_model(mdl, lora_cfg)

# 4) Formateo con plantilla de chat oficial de Qwen2.5
def format_example(example):
    text = tok.apply_chat_template(
        example["messages"],
        tokenize=False,
        add_generation_prompt=False,
    )
    return {"text": text + tok.eos_token}

train_ds_fmt = train_ds.map(format_example, remove_columns=train_ds.column_names)
eval_ds_fmt = eval_ds.map(format_example, remove_columns=eval_ds.column_names)

# 5) SFTTrainer
output_dir = "/content/drive/MyDrive/ASISTENTE-SFTT-Qwen2p5-7B-LoRA"

args = TrainingArguments(
    output_dir=output_dir,
    num_train_epochs=2,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    gradient_accumulation_steps=4,
    learning_rate=2e-4,
    logging_steps=10,
    eval_strategy="steps",
    eval_steps=100,
    save_steps=100,
    save_total_limit=2,
    do_eval=True,
    bf16=True,
    report_to="wandb",
    run_name="sftt-qwen2p5",
)

trainer = SFTTrainer(
    model=mdl,
    args=args,
    train_dataset=train_ds_fmt,
    eval_dataset=eval_ds_fmt,
)

trainer.train()
trainer.save_model(output_dir)

# 6) Fusión de adaptadores a modelo base (bf16, sin cuantizar)
base_merge = AutoModelForCausalLM.from_pretrained(
    base_model,
    device_map="auto",
    torch_dtype=torch.bfloat16,
)

model_with_lora = PeftModel.from_pretrained(base_merge, output_dir)
merged_dir = "/content/drive/MyDrive/ASISTENTE-SFTT-Qwen2p5-7B-LoRA-merged"
merged_model = model_with_lora.merge_and_unload()

merged_model.save_pretrained(merged_dir, max_shard_size="2GB")
tok.save_pretrained(merged_dir)