Update app.py

app.py

@@ -1,188 +1,186 @@
 """
-PromptWizard Qwen Training — Gita Edition (Fixed Tokenizer Import)
-Fine-tunes Qwen using rahul7star/Gita dataset (.csv)
-Uploads trained model to rahul7star/Qwen0.5-3B-Gita
+PromptWizard — Qwen2.5-0.5B Fine-tuning on Bhagavad Gita Dataset
+Optimized for Hugging Face Spaces with Zero GPU Support
 """
 
 import gradio as gr
 import spaces
 import torch
+import pandas as pd
+from datasets import Dataset
 from transformers import (
-    AutoModelForCausalLM,
     AutoTokenizer,
+    AutoModelForCausalLM,
     Trainer,
     TrainingArguments,
 )
-from datasets import load_dataset
 from peft import LoraConfig, get_peft_model, TaskType
-from huggingface_hub import HfApi, HfFolder, Repository
-import os, tempfile, shutil
+from huggingface_hub import HfApi
+import os
 
 
-# === GPU check ===
+# ------------------------------------------------------
+# 🧠 Helper: Check GPU availability
+# ------------------------------------------------------
 def check_gpu_status():
-    return "🚀 Zero GPU Ready - GPU will be allocated when training starts"
+    return "🟢 Ready — GPU will be auto-assigned when training starts."
 
 
-# === Training function ===
-@spaces.GPU(duration=300)
+# ------------------------------------------------------
+# 🚀 Main Training Function (runs on GPU)
+# ------------------------------------------------------
+@spaces.GPU(duration=600)  # 10 minutes GPU allocation
 def train_model(model_name, num_epochs, batch_size, learning_rate, progress=gr.Progress()):
-    progress(0, desc="Initializing...")
-    output_log = []
-
+    log = []
     try:
-        # ==== Device ====
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        output_log.append(f"🎮 Using device: {device}")
-        if device == "cuda":
-            output_log.append(f"✅ GPU: {torch.cuda.get_device_name(0)}")
-
-        # ==== Load dataset ====
-        progress(0.1, desc="Loading rahul7star/Gita dataset...")
-        output_log.append("\n📚 Loading dataset from rahul7star/Gita...")
-
-        dataset = load_dataset("rahul7star/Gita", split="train")
-        output_log.append(f"   Loaded {len(dataset)} samples")
-        output_log.append(f"   Columns: {dataset.column_names}")
-
-        # ==== Format dataset ====
-        def format_example(item):
-            text = (
-                item.get("text")
-                or item.get("content")
-                or " ".join(str(v) for v in item.values())
-            )
-            prompt = f"""<|system|>
-You are a wise teacher interpreting Bhagavad Gita with deep insights.
-<|user|>
-{text}
-<|assistant|>
-"""
-            return {"text": prompt}
+        log.append(f"🎮 Device: {device}")
 
-        dataset = dataset.map(format_example)
-        output_log.append(f"   ✅ Formatted {len(dataset)} examples")
+        # ------------------------------------------------------
+        # 📂 Load Dataset from Hugging Face Repo (rahul7star/Gita)
+        # ------------------------------------------------------
+        progress(0.1, desc="Loading Bhagavad Gita dataset...")
+        log.append("\n📚 Loading dataset from rahul7star/Gita ...")
+
+        # Load CSV from your repo
+        dataset = Dataset.from_pandas(
+            pd.read_csv("https://huggingface.co/datasets/rahul7star/Gita/resolve/main/Gita.csv")
+        )
 
-        # ==== Model & Tokenizer ====
-        # Load model and tokenizer
+        # 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'")
+
+        def format_row(row):
+            return {
+                "text": f"<|system|>\nYou are a spiritual teacher summarizing Gita verses.\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")
+
+        # ------------------------------------------------------
+        # 🤖 Load Qwen Model and Tokenizer
+        # ------------------------------------------------------
         progress(0.3, desc="Loading model and tokenizer...")
-        output_log.append(f"\n🤖 Loading {model_name}...")
+        log.append(f"\n🔍 Loading model: {model_name}")
 
-# Always use smaller model for Zero GPU demo
-        model_name = "Qwen/Qwen2.5-0.5B"  # Use smallest model for Zero GPU
-        output_log.append("   Note: Using 0.5B model for Zero GPU compatibility")
+        model_name = "Qwen/Qwen2.5-0.5B"  # safest base model for Zero GPU
 
-# ✅ Correct tokenizer and model loading
         tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-# Fix for padding token (Qwen sometimes lacks it)
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
 
         model = AutoModelForCausalLM.from_pretrained(
-           model_name,
-           torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+            model_name,
+            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
         )
 
-# Move to GPU if available
-if device == "cuda":
-    model = model.to(device)
+        if device == "cuda":
+            model = model.to(device)
 
-output_log.append("   Model loaded successfully with correct tokenizer and model setup")
-        # ==== LoRA ====
-        progress(0.4, desc="Configuring LoRA...")
-        output_log.append("\n⚙️ Setting up LoRA for efficient fine-tuning...")
+        log.append("✅ Model & tokenizer loaded successfully")
 
+        # ------------------------------------------------------
+        # ⚙️ Configure LoRA for lightweight fine-tuning
+        # ------------------------------------------------------
+        progress(0.4, desc="Configuring LoRA...")
         lora_config = LoraConfig(
             task_type=TaskType.CAUSAL_LM,
             r=8,
             lora_alpha=16,
             lora_dropout=0.1,
             target_modules=["q_proj", "v_proj"],
-            bias="none",
         )
         model = get_peft_model(model, lora_config)
 
         trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
-        output_log.append(f"   Trainable parameters: {trainable_params:,}")
+        total_params = sum(p.numel() for p in model.parameters())
+        log.append(f"🧩 Trainable params: {trainable_params:,} / {total_params:,}")
 
-        # ==== Tokenization ====
+        # ------------------------------------------------------
+        # ✂️ Tokenize Dataset
+        # ------------------------------------------------------
         progress(0.5, desc="Tokenizing dataset...")
+
         def tokenize_fn(examples):
             return tokenizer(
                 examples["text"],
                 padding="max_length",
                 truncation=True,
-                max_length=256,
+                max_length=512,
             )
 
-        dataset = dataset.map(tokenize_fn, batched=True)
-        output_log.append("   ✅ Tokenization done")
+        tokenized = dataset.map(tokenize_fn, batched=True)
 
-        # ==== Training setup ====
-        progress(0.6, desc="Setting up Trainer...")
-        output_dir = "./qwen-gita-lora"
+        # ------------------------------------------------------
+        # 🎯 Setup Training Arguments
+        # ------------------------------------------------------
+        progress(0.6, desc="Configuring training...")
         training_args = TrainingArguments(
-            output_dir=output_dir,
+            output_dir="./qwen-gita-output",
             num_train_epochs=num_epochs,
             per_device_train_batch_size=batch_size,
             gradient_accumulation_steps=2,
-            warmup_steps=10,
+            learning_rate=learning_rate,
+            warmup_steps=5,
             logging_steps=5,
-            save_strategy="epoch",
+            save_strategy="no",
             fp16=device == "cuda",
-            optim="adamw_torch",
-            learning_rate=learning_rate,
-            max_steps=100,
+            max_steps=100,  # short demo run
+            report_to="none",
         )
 
+        # ------------------------------------------------------
+        # 🏋️ Train
+        # ------------------------------------------------------
+        progress(0.7, desc="Training model...")
+        log.append("\n🚀 Starting training on Bhagavad Gita dataset...")
+
         trainer = Trainer(
             model=model,
             args=training_args,
-            train_dataset=dataset,
+            train_dataset=tokenized,
             tokenizer=tokenizer,
         )
+        trainer.train()
 
-        # ==== Train ====
-        progress(0.7, desc="Training...")
-        output_log.append("\n🚀 Starting training...\n" + "=" * 50)
-        train_result = trainer.train()
+        progress(0.9, desc="Finalizing and saving...")
 
-        progress(0.85, desc="Saving model...")
-        output_log.append("\n💾 Saving model locally...")
-        trainer.save_model(output_dir)
+        # ------------------------------------------------------
+        # 💾 Save and Upload to HF Repo
+        # ------------------------------------------------------
+        output_dir = "./qwen-gita-lora"
+        os.makedirs(output_dir, exist_ok=True)
+        model.save_pretrained(output_dir)
         tokenizer.save_pretrained(output_dir)
 
-        # ==== Upload to HF Hub ====
-        progress(0.9, desc="Uploading to Hugging Face Hub...")
-        hf_repo = "rahul7star/Qwen0.5-3B-Gita"
-        output_log.append(f"\n☁️ Uploading fine-tuned model to: {hf_repo}")
-
+        log.append("\n📤 Uploading model to 🤗 Hugging Face Hub...")
         api = HfApi()
-        token = HfFolder.get_token()
-        api.create_repo(repo_id=hf_repo, exist_ok=True)
-
-        with tempfile.TemporaryDirectory() as tmpdir:
-            repo = Repository(local_dir=tmpdir, clone_from=hf_repo, use_auth_token=token)
-            shutil.copytree(output_dir, tmpdir, dirs_exist_ok=True)
-            repo.push_to_hub(commit_message="Upload fine-tuned Qwen-Gita model")
+        repo_id = "rahul7star/Qwen0.5-3B-Gita"
+        api.upload_folder(folder_path=output_dir, repo_id=repo_id)
+        log.append(f"✅ Uploaded LoRA fine-tuned model to {repo_id}")
 
         progress(1.0, desc="Complete!")
-        output_log.append("\n✅ Training complete & model uploaded successfully!")
+        log.append("\n🎉 Training finished successfully!")
 
     except Exception as e:
-        output_log.append(f"\n❌ Error: {e}")
+        log.append(f"\n❌ Error: {e}")
 
-    return "\n".join(output_log)
+    return "\n".join(log)
 
 
-# === Gradio UI ===
+# ------------------------------------------------------
+# 🎨 Gradio Interface
+# ------------------------------------------------------
 def create_interface():
     with gr.Blocks(title="PromptWizard — Qwen Gita Trainer") as demo:
         gr.Markdown("""
-        # 🧘 Qwen Fine-tuning — Gita Edition  
-        Fine-tune **Qwen 0.5B** using your [rahul7star/Gita](https://huggingface.co/datasets/rahul7star/Gita) dataset  
-        and automatically upload to your repo **rahul7star/Qwen0.5-3B-Gita**.
+        # 🧘 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**
         """)
 
         with gr.Row():
@@ -190,38 +188,45 @@ def create_interface():
                 gpu_status = gr.Textbox(
                     label="GPU Status",
                     value=check_gpu_status(),
-                    interactive=False,
+                    interactive=False
                 )
-
-                # hidden constant input for model name
                 model_name = gr.Textbox(
                     value="Qwen/Qwen2.5-0.5B",
-                    label="Base Model",
-                    interactive=False,
                     visible=False
                 )
-
-                num_epochs = gr.Slider(1, 3, value=1, step=1, label="Epochs")
-                batch_size = gr.Slider(1, 4, value=2, step=1, label="Batch Size")
+                num_epochs = gr.Slider(1, 3, 1, step=1, label="Epochs")
+                batch_size = gr.Slider(1, 4, 2, step=1, label="Batch Size")
                 learning_rate = gr.Number(value=5e-5, label="Learning Rate")
                 train_btn = gr.Button("🚀 Start Fine-tuning", variant="primary")
 
             with gr.Column():
                 output = gr.Textbox(
-                    label="Training Log",
+                    label="Training Logs",
                     lines=25,
                     max_lines=40,
-                    value="Click 'Start Fine-tuning' to train on the Gita dataset and upload to your model repo.",
+                    value="Press Start to fine-tune Qwen on your Gita dataset.\nZero GPU will allocate automatically.",
                 )
 
-        # ✅ Fixed: pass components, not strings
         train_btn.click(
             fn=train_model,
             inputs=[model_name, num_epochs, batch_size, learning_rate],
             outputs=output,
         )
 
+        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`
+        """)
+
     return demo
+
+
+# ------------------------------------------------------
+# 🚪 Launch App
+# ------------------------------------------------------
 if __name__ == "__main__":
     demo = create_interface()
     demo.launch()