import unsloth  # must be first
import pandas as pd
import torch
from datasets import Dataset
from transformers import TrainingArguments
from unsloth import FastLanguageModel
from trl import SFTTrainer
from shutil import copytree
import os

# Load and format your dataset
df = pd.read_csv("data.csv")
df["text"] = df.apply(lambda row: f"### Instruction:\n{row['instruction']}\n\n### Response:\n{row['response']}\n", axis=1)
dataset = Dataset.from_pandas(df[["text"]])

# Load Unsloth model
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/Llama-3.2-3B-Instruct",  # or another Unsloth-compatible model
    max_seq_length = 2048,
    dtype = torch.float16,
    load_in_4bit = True,
)

# Apply LoRA without task_type
model = FastLanguageModel.get_peft_model(
    model,
    r = 8,
    lora_alpha = 32,
    lora_dropout = 0.05,
    bias = "none",
)

# Tokenize
def tokenize(example):
    return tokenizer(example["text"], truncation=True, padding="max_length", max_length=512)

tokenized_dataset = dataset.map(tokenize, batched=True)

# Define training args
training_args = TrainingArguments(
    output_dir = "./lora-finetuned",
    per_device_train_batch_size = 2,
    num_train_epochs = 3,
    learning_rate = 2e-4,
    logging_steps = 10,
    save_steps = 100,
    fp16 = True,
)

# Train
trainer = SFTTrainer(
    model = model,
    tokenizer = tokenizer,
    args = training_args,
    train_dataset = tokenized_dataset,
)

trainer.train()

# Save model
model.save_pretrained("./lora-finetuned")

# Copy to visible folder
try:
    copytree("./lora-finetuned", "/home/user/app/final_model", 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))