Update train.py

train.py

@@ -1,87 +1,68 @@
-# === Required early import ===
-import unsloth
+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 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 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 model ===
+# Load Unsloth model
 model, tokenizer = FastLanguageModel.from_pretrained(
-    model_name = MODEL_NAME,
-    max_seq_length = MAX_SEQ_LENGTH,
+    model_name = "unsloth/Llama-3.2-3B-Instruct",  # or another Unsloth-compatible model
+    max_seq_length = 2048,
     dtype = torch.float16,
-    load_in_4bit = USE_4BIT,
+    load_in_4bit = True,
 )
 
-# === 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
+# Apply LoRA without task_type
+model = FastLanguageModel.get_peft_model(
+    model,
+    r = 8,
+    lora_alpha = 32,
+    lora_dropout = 0.05,
+    bias = "none",
+)
 
-dataset = dataset.map(add_text_column)
+# Tokenize
+def tokenize(example):
+    return tokenizer(example["text"], truncation=True, padding="max_length", max_length=512)
 
-# === Formatting function for trainer ===
-def formatting_func(example):
-    return [example["text"]]
+tokenized_dataset = dataset.map(tokenize, batched=True)
 
-# === TrainingArguments ===
+# Define training args
 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",
+    output_dir = "./lora-finetuned",
+    per_device_train_batch_size = 2,
+    num_train_epochs = 3,
+    learning_rate = 2e-4,
     logging_steps = 10,
-    save_strategy = "epoch",
-    learning_rate = LR,
-    bf16 = False,
+    save_steps = 100,
     fp16 = True,
-    max_grad_norm = 1.0,
-    warmup_ratio = 0.03,
-    lr_scheduler_type = "linear",
-    disable_tqdm = False,
-    report_to = "none",
 )
 
-# === Trainer setup ===
+# Train
 trainer = SFTTrainer(
     model = model,
     tokenizer = tokenizer,
-    train_dataset = dataset,
-    dataset_text_field = "text",
-    formatting_func = formatting_func,
     args = training_args,
+    train_dataset = tokenized_dataset,
 )
 
-# === 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)
+# Save model
+model.save_pretrained("./lora-finetuned")
 
-# === Copy to visible directory ===
+# Copy to visible folder
 try:
-    copytree(save_dir, final_dir, dirs_exist_ok=True)
+    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))