# === Required early import ===
import unsloth
from unsloth import FastLanguageModel
from trl import SFTTrainer
from transformers import TrainingArguments
from datasets import load_dataset
import torch
from shutil import copytree
import os

# === Model and training config ===
MODEL_NAME = "unsloth/SmolLM2-1.7B-Instruct"  # Change if using another model
BATCH_SIZE = 2
EPOCHS = 3
LR = 2e-4
MAX_SEQ_LENGTH = 2048
USE_4BIT = True

# === Load model ===
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = MODEL_NAME,
    max_seq_length = MAX_SEQ_LENGTH,
    dtype = torch.float16,
    load_in_4bit = USE_4BIT,
)

# === Load dataset ===
dataset = load_dataset("csv", data_files="data.csv")["train"]  # Replace with your CSV if needed

# === Create 'text' column from instruction + response ===
def add_text_column(example):
    example["text"] = example["instruction"] + "\n" + example["response"]
    return example

dataset = dataset.map(add_text_column)

# === Formatting function for trainer ===
def formatting_func(example):
    return [example["text"]]

# === TrainingArguments ===
training_args = TrainingArguments(
    output_dir = "output",
    num_train_epochs = EPOCHS,
    per_device_train_batch_size = BATCH_SIZE,
    gradient_accumulation_steps = 1,
    gradient_checkpointing = True,
    optim = "paged_adamw_8bit",
    logging_steps = 10,
    save_strategy = "epoch",
    learning_rate = LR,
    bf16 = False,
    fp16 = True,
    max_grad_norm = 1.0,
    warmup_ratio = 0.03,
    lr_scheduler_type = "linear",
    disable_tqdm = False,
    report_to = "none",
)

# === Trainer setup ===
trainer = SFTTrainer(
    model = model,
    tokenizer = tokenizer,
    train_dataset = dataset,
    dataset_text_field = "text",
    formatting_func = formatting_func,
    args = training_args,
)

# === Train ===
model = FastLanguageModel.prepare_for_training(model)
trainer.train()

# === Save model ===
save_dir = "output"
final_dir = "/home/user/app/final_model"

model.save_pretrained(save_dir, safe_serialization=True)
tokenizer.save_pretrained(save_dir)

# === Copy to visible directory ===
try:
    copytree(save_dir, final_dir, dirs_exist_ok=True)
    print("✅ Model saved to /home/user/app/final_model for download in UI.")
except Exception as e:
    print("⚠️ Failed to copy model to visible folder:", str(e))