Update train.py

train.py

@@ -1,12 +1,14 @@
+# ✅ Final train.py with ZIP logic added
 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  # ✅ now works because we added 'trl'
+from trl import SFTTrainer
 import os
 import shutil
+import zipfile
 
 # Load and format your dataset
 df = pd.read_csv("data.csv")
@@ -38,7 +40,7 @@ tokenized_dataset = dataset.map(tokenize, batched=True)
 
 # Set up training
 training_args = TrainingArguments(
-    output_dir = "./lora-finetuned",
+    output_dir = "./output_model",
     per_device_train_batch_size = 2,
     num_train_epochs = 3,
     learning_rate = 2e-4,
@@ -57,15 +59,20 @@ trainer = SFTTrainer(
 
 trainer.train()
 
-# ✅ Save to Hugging Face persistent path
-output_path = "./outputs/final_model"
-os.makedirs(output_path, exist_ok=True)
+# Save the fine-tuned LoRA adapter
+output_dir = "./output_model"
+os.makedirs(output_dir, exist_ok=True)
+model.save_pretrained(output_dir)
 
+# ✅ Zip it for download
+zip_path = "/home/user/app/model.zip"
 try:
-    model.save_pretrained(output_path)
-    print(f"✅ Model saved to {output_path}")
-    print("📁 Contents of saved model directory:")
-    for file in os.listdir(output_path):
-        print(" -", file)
+    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
+        for root, _, files in os.walk(output_dir):
+            for file in files:
+                full_path = os.path.join(root, file)
+                rel_path = os.path.relpath(full_path, output_dir)
+                zipf.write(full_path, rel_path)
+    print(f"✅ Zipped model to {zip_path}")
 except Exception as e:
-    print(f"❌ Failed to save model: {e}")
+    print(f"❌ Failed to zip model: {e}")