Update app.py

app.py

@@ -1,6 +1,6 @@
 """
 PromptWizard — Qwen2.5-0.5B Fine-tuning on Bhagavad Gita Dataset
-Optimized for Hugging Face Spaces with Zero GPU Support
+Downloads CSV locally before training (for Hugging Face Spaces)
 """
 
 import gradio as gr
@@ -15,59 +15,90 @@ from transformers import (
     TrainingArguments,
 )
 from peft import LoraConfig, get_peft_model, TaskType
-from huggingface_hub import HfApi
+from huggingface_hub import snapshot_download, HfApi
 import os
+import shutil
 
 
 # ------------------------------------------------------
-# 🧠 Helper: Check GPU availability
+# 🧠 GPU check
 # ------------------------------------------------------
 def check_gpu_status():
-    return "🟢 Ready — GPU will be auto-assigned when training starts."
+    return "🟢 Ready — GPU will be assigned at runtime (Zero GPU mode)"
 
 
 # ------------------------------------------------------
-# 🚀 Main Training Function (runs on GPU)
+# 🧩 Download Dataset to /tmp/
 # ------------------------------------------------------
-@spaces.GPU(duration=600)  # 10 minutes GPU allocation
+def download_gita_dataset():
+    repo_id = "rahul7star/Gita"
+    local_dir = "/tmp/gita_data"
+
+    if os.path.exists(local_dir):
+        shutil.rmtree(local_dir)
+    os.makedirs(local_dir, exist_ok=True)
+
+    print(f"📥 Downloading dataset from {repo_id} ...")
+    snapshot_download(repo_id=repo_id, local_dir=local_dir, repo_type="dataset")
+
+    # Try to locate the CSV file
+    csv_path = None
+    for root, _, files in os.walk(local_dir):
+        for f in files:
+            if f.lower().endswith(".csv"):
+                csv_path = os.path.join(root, f)
+                break
+    if not csv_path:
+        raise FileNotFoundError("No CSV file found in the Gita dataset repository.")
+
+    print(f"✅ Found CSV: {csv_path}")
+    return csv_path
+
+
+# ------------------------------------------------------
+# 🚀 Training function
+# ------------------------------------------------------
+@spaces.GPU(duration=300)
 def train_model(model_name, num_epochs, batch_size, learning_rate, progress=gr.Progress()):
-    log = []
+    logs = []
     try:
+        progress(0.05, desc="Initializing...")
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        log.append(f"🎮 Device: {device}")
+        logs.append(f"🎮 Device: {device}")
 
         # ------------------------------------------------------
-        # 📂 Load Dataset from Hugging Face Repo (rahul7star/Gita)
+        # 📂 Step 1: Download dataset
         # ------------------------------------------------------
-        progress(0.1, desc="Loading Bhagavad Gita dataset...")
-        log.append("\n📚 Loading dataset from rahul7star/Gita ...")
+        progress(0.1, desc="Downloading dataset...")
+        logs.append("\n📥 Downloading Gita dataset from HF Hub...")
+        csv_path = download_gita_dataset()
 
-        # Load CSV from your repo
-        dataset = Dataset.from_pandas(
-            pd.read_csv("https://huggingface.co/datasets/rahul7star/Gita/resolve/main/Gita.csv")
-        )
+        # ------------------------------------------------------
+        # 🧾 Step 2: Load dataset from CSV
+        # ------------------------------------------------------
+        progress(0.2, desc="Loading dataset...")
+        df = pd.read_csv(csv_path)
+        if not all(c in df.columns for c in ["question", "answer"]):
+            raise ValueError("CSV must contain 'question' and 'answer' columns.")
 
-        # Expected columns: 'question', 'answer'
-        if not all(c in dataset.column_names for c in ["question", "answer"]):
-            raise ValueError("CSV must contain columns: 'question' and 'answer'")
+        hf_dataset = Dataset.from_pandas(df)
 
         def format_row(row):
             return {
-                "text": f"<|system|>\nYou are a spiritual teacher summarizing Gita verses.\n"
+                "text": f"<|system|>\nYou are a spiritual guide explaining Gita concepts.\n"
                         f"<|user|>\n{row['question']}\n"
                         f"<|assistant|>\n{row['answer']}"
             }
 
-        dataset = dataset.map(format_row)
-        log.append(f"✅ Loaded {len(dataset)} examples from Bhagavad Gita CSV")
+        hf_dataset = hf_dataset.map(format_row)
+        logs.append(f"✅ Loaded {len(hf_dataset)} examples from {csv_path}")
 
         # ------------------------------------------------------
-        # 🤖 Load Qwen Model and Tokenizer
+        # 🤖 Step 3: Load model + tokenizer
         # ------------------------------------------------------
-        progress(0.3, desc="Loading model and tokenizer...")
-        log.append(f"\n🔍 Loading model: {model_name}")
-
-        model_name = "Qwen/Qwen2.5-0.5B"  # safest base model for Zero GPU
+        progress(0.3, desc="Loading Qwen model...")
+        model_name = "Qwen/Qwen2.5-0.5B"
+        logs.append(f"\n🔍 Loading base model: {model_name}")
 
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         if tokenizer.pad_token is None:
@@ -77,14 +108,12 @@ def train_model(model_name, num_epochs, batch_size, learning_rate, progress=gr.P
             model_name,
             torch_dtype=torch.float16 if device == "cuda" else torch.float32,
         )
-
         if device == "cuda":
             model = model.to(device)
-
-        log.append("✅ Model & tokenizer loaded successfully")
+        logs.append("✅ Model and tokenizer ready")
 
         # ------------------------------------------------------
-        # ⚙️ Configure LoRA for lightweight fine-tuning
+        # ⚙️ Step 4: Apply LoRA config
         # ------------------------------------------------------
         progress(0.4, desc="Configuring LoRA...")
         lora_config = LoraConfig(
@@ -96,31 +125,28 @@ def train_model(model_name, num_epochs, batch_size, learning_rate, progress=gr.P
         )
         model = get_peft_model(model, lora_config)
 
-        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
-        total_params = sum(p.numel() for p in model.parameters())
-        log.append(f"🧩 Trainable params: {trainable_params:,} / {total_params:,}")
-
         # ------------------------------------------------------
-        # ✂️ Tokenize Dataset
+        # ✂️ Step 5: Tokenize dataset
         # ------------------------------------------------------
-        progress(0.5, desc="Tokenizing dataset...")
+        progress(0.5, desc="Tokenizing data...")
 
-        def tokenize_fn(examples):
+        def tokenize_fn(batch):
             return tokenizer(
-                examples["text"],
-                padding="max_length",
+                batch["text"],
                 truncation=True,
+                padding="max_length",
                 max_length=512,
             )
 
-        tokenized = dataset.map(tokenize_fn, batched=True)
+        tokenized = hf_dataset.map(tokenize_fn, batched=True)
+        logs.append("🧾 Dataset tokenized successfully")
 
         # ------------------------------------------------------
-        # 🎯 Setup Training Arguments
+        # 🎯 Step 6: Training arguments
         # ------------------------------------------------------
         progress(0.6, desc="Configuring training...")
         training_args = TrainingArguments(
-            output_dir="./qwen-gita-output",
+            output_dir="/tmp/qwen-gita-output",
             num_train_epochs=num_epochs,
             per_device_train_batch_size=batch_size,
             gradient_accumulation_steps=2,
@@ -129,15 +155,15 @@ def train_model(model_name, num_epochs, batch_size, learning_rate, progress=gr.P
             logging_steps=5,
             save_strategy="no",
             fp16=device == "cuda",
-            max_steps=100,  # short demo run
+            max_steps=100,
             report_to="none",
         )
 
         # ------------------------------------------------------
-        # 🏋️ Train
+        # 🏋️ Step 7: Train model
         # ------------------------------------------------------
-        progress(0.7, desc="Training model...")
-        log.append("\n🚀 Starting training on Bhagavad Gita dataset...")
+        progress(0.7, desc="Training in progress...")
+        logs.append("\n🚀 Starting fine-tuning...")
 
         trainer = Trainer(
             model=model,
@@ -147,29 +173,28 @@ def train_model(model_name, num_epochs, batch_size, learning_rate, progress=gr.P
         )
         trainer.train()
 
-        progress(0.9, desc="Finalizing and saving...")
-
         # ------------------------------------------------------
-        # 💾 Save and Upload to HF Repo
+        # 💾 Step 8: Save + Upload
         # ------------------------------------------------------
-        output_dir = "./qwen-gita-lora"
+        progress(0.9, desc="Saving and uploading...")
+        output_dir = "/tmp/qwen-gita-lora"
         os.makedirs(output_dir, exist_ok=True)
         model.save_pretrained(output_dir)
         tokenizer.save_pretrained(output_dir)
 
-        log.append("\n📤 Uploading model to 🤗 Hugging Face Hub...")
-        api = HfApi()
+        logs.append("\n📤 Uploading fine-tuned LoRA model to Hugging Face Hub...")
         repo_id = "rahul7star/Qwen0.5-3B-Gita"
+        api = HfApi()
         api.upload_folder(folder_path=output_dir, repo_id=repo_id)
-        log.append(f"✅ Uploaded LoRA fine-tuned model to {repo_id}")
+        logs.append(f"✅ Uploaded fine-tuned model to {repo_id}")
 
         progress(1.0, desc="Complete!")
-        log.append("\n🎉 Training finished successfully!")
+        logs.append("\n🎉 Training complete!")
 
     except Exception as e:
-        log.append(f"\n❌ Error: {e}")
+        logs.append(f"\n❌ Error: {str(e)}")
 
-    return "\n".join(log)
+    return "\n".join(logs)
 
 
 # ------------------------------------------------------
@@ -178,9 +203,9 @@ def train_model(model_name, num_epochs, batch_size, learning_rate, progress=gr.P
 def create_interface():
     with gr.Blocks(title="PromptWizard — Qwen Gita Trainer") as demo:
         gr.Markdown("""
-        # 🧘 Qwen Gita Trainer — Fine-tune Qwen 0.5B  
-        Train Qwen2.5-0.5B on your **Bhagavad Gita dataset (CSV)**  
-        and auto-upload LoRA weights to your repo: **rahul7star/Qwen0.5-3B-Gita**
+        # 🧘 PromptWizard — Qwen2.5-0.5B Gita Trainer  
+        Fine-tune Qwen 0.5B on your **Bhagavad Gita CSV dataset**  
+        Automatically uploads LoRA weights to `rahul7star/Qwen0.5-3B-Gita`
         """)
 
         with gr.Row():
@@ -188,11 +213,11 @@ def create_interface():
                 gpu_status = gr.Textbox(
                     label="GPU Status",
                     value=check_gpu_status(),
-                    interactive=False
+                    interactive=False,
                 )
                 model_name = gr.Textbox(
                     value="Qwen/Qwen2.5-0.5B",
-                    visible=False
+                    visible=False,
                 )
                 num_epochs = gr.Slider(1, 3, 1, step=1, label="Epochs")
                 batch_size = gr.Slider(1, 4, 2, step=1, label="Batch Size")
@@ -204,7 +229,7 @@ def create_interface():
                     label="Training Logs",
                     lines=25,
                     max_lines=40,
-                    value="Press Start to fine-tune Qwen on your Gita dataset.\nZero GPU will allocate automatically.",
+                    value="Click 'Start Fine-tuning' to train on Bhagavad Gita dataset...",
                 )
 
         train_btn.click(
@@ -215,17 +240,17 @@ def create_interface():
 
         gr.Markdown("""
         ---
-        💡 **Notes**  
-        - Requires your CSV on HF dataset: `rahul7star/Gita/Gita.csv`  
-        - Columns: `question`, `answer`  
-        - Model uploads automatically to: `rahul7star/Qwen0.5-3B-Gita`
+        **Notes:**
+        - Downloads dataset: `rahul7star/Gita` → `/tmp/gita_data/Gita.csv`
+        - Trains using LoRA for efficiency
+        - Uploads to `rahul7star/Qwen0.5-3B-Gita`
         """)
 
     return demo
 
 
 # ------------------------------------------------------
-# 🚪 Launch App
+# 🚪 Launch app
 # ------------------------------------------------------
 if __name__ == "__main__":
     demo = create_interface()