Update app.py

app.py

@@ -1,236 +1,150 @@
-"""
-PromptWizard — Qwen2.5-0.5B Fine-tuning on Bhagavad Gita Dataset
-Downloads CSV locally before training (for Hugging Face Spaces)
-"""
-
-import gradio as gr
-import spaces
+import os
 import torch
-import pandas as pd
-from datasets import Dataset
+import gradio as gr
+from datasets import load_dataset
 from transformers import (
     AutoTokenizer,
     AutoModelForCausalLM,
     Trainer,
     TrainingArguments,
+    DataCollatorForLanguageModeling,
 )
-from peft import LoraConfig, get_peft_model, TaskType
-from huggingface_hub import snapshot_download, HfApi
-import os
-import shutil
+from huggingface_hub import HfApi, HfFolder
 
+# ---------------------------------------------------------------------
+# GPU check
+# ---------------------------------------------------------------------
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# ------------------------------------------------------
-# 🧠 GPU check
-# ------------------------------------------------------
 def check_gpu_status():
-    return "🟢 Ready — GPU will be assigned at runtime (Zero GPU mode)"
-
-
-# ------------------------------------------------------
-# 🧩 Download Dataset to /tmp/
-# ------------------------------------------------------
-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()):
-    logs = []
-    try:
-        progress(0.05, desc="Initializing...")
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        logs.append(f"🎮 Device: {device}")
-
-        # ------------------------------------------------------
-        # 📂 Step 1: Download dataset
-        # ------------------------------------------------------
-        progress(0.1, desc="Downloading dataset...")
-        logs.append("\n📥 Downloading Gita dataset from HF Hub...")
-        csv_path = download_gita_dataset()
-
-        # ------------------------------------------------------
-        # 🧾 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.")
-
-        hf_dataset = Dataset.from_pandas(df)
-
-        def format_row(row):
-            return {
-                "text": f"<|system|>\nYou are a spiritual guide explaining Gita concepts.\n"
-                        f"<|user|>\n{row['question']}\n"
-                        f"<|assistant|>\n{row['answer']}"
-            }
-
-        hf_dataset = hf_dataset.map(format_row)
-        logs.append(f"✅ Loaded {len(hf_dataset)} examples from {csv_path}")
-
-        # ------------------------------------------------------
-        # 🤖 Step 3: Load model + tokenizer
-        # ------------------------------------------------------
-        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:
-            tokenizer.pad_token = tokenizer.eos_token
-
-        model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+    return f"✅ GPU: {torch.cuda.get_device_name(0)}" if device == "cuda" else "⚠️ Using CPU only"
+
+# ---------------------------------------------------------------------
+# Training Logic
+# ---------------------------------------------------------------------
+def train_model(model_name, num_epochs, batch_size, learning_rate, progress=gr.Progress(track_tqdm=True)):
+    output_log = []
+
+    # ==== 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 item.get("verse")
+            or " ".join(str(v) for v in item.values())
         )
-        if device == "cuda":
-            model = model.to(device)
-        logs.append("✅ Model and tokenizer ready")
+        prompt = f"""<|system|>
+You are a wise teacher interpreting Bhagavad Gita with deep insights.
+<|user|>
+{text}
+<|assistant|>
+"""
+        return {"text": prompt}
 
-        # ------------------------------------------------------
-        # ⚙️ Step 4: Apply LoRA config
-        # ------------------------------------------------------
-        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"],
-        )
-        model = get_peft_model(model, lora_config)
+    dataset = dataset.map(format_example)
+    output_log.append(f"   ✅ Formatted {len(dataset)} examples")
 
-        # ------------------------------------------------------
-        # ✂️ Step 5: Tokenize dataset
-        # ------------------------------------------------------
-        progress(0.5, desc="Tokenizing data...")
+    # ==== Load tokenizer & model ====
+    progress(0.3, desc="Loading model and tokenizer...")
+    output_log.append("\n🤖 Loading Qwen model and tokenizer...")
 
-        def tokenize_fn(batch):
-            return tokenizer(
-                batch["text"],
-                truncation=True,
-                padding="max_length",
-                max_length=512,
-            )
+    base_model = "Qwen/Qwen2.5-0.5B"
+    tokenizer = AutoTokenizer.from_pretrained(base_model)
 
-        tokenized = hf_dataset.map(tokenize_fn, batched=True)
-        logs.append("🧾 Dataset tokenized successfully")
+    # Fix missing pad token
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
 
-        # ------------------------------------------------------
-        # 🎯 Step 6: Training arguments
-        # ------------------------------------------------------
-        progress(0.6, desc="Configuring training...")
-        training_args = TrainingArguments(
-            output_dir="/tmp/qwen-gita-output",
-            num_train_epochs=num_epochs,
-            per_device_train_batch_size=batch_size,
-            gradient_accumulation_steps=2,
-            learning_rate=learning_rate,
-            warmup_steps=5,
-            logging_steps=5,
-            save_strategy="no",
-            fp16=device == "cuda",
-            max_steps=100,
-            report_to="none",
-        )
+    model = AutoModelForCausalLM.from_pretrained(
+        base_model,
+        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+    ).to(device)
 
-        # ------------------------------------------------------
-        # 🏋️ Step 7: Train model
-        # ------------------------------------------------------
-        progress(0.7, desc="Training in progress...")
-        logs.append("\n🚀 Starting fine-tuning...")
+    # ==== Tokenize dataset ====
+    progress(0.4, desc="Tokenizing dataset...")
+    output_log.append("\n✏️ Tokenizing dataset...")
 
-        trainer = Trainer(
-            model=model,
-            args=training_args,
-            train_dataset=tokenized,
-            tokenizer=tokenizer,
+    def tokenize_function(examples):
+        return tokenizer(
+            examples["text"],
+            truncation=True,
+            padding="max_length",
+            max_length=512,
         )
-        trainer.train()
-
-        # ------------------------------------------------------
-        # 💾 Step 8: Save + Upload
-        # ------------------------------------------------------
-        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)
 
-        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)
-        logs.append(f"✅ Uploaded fine-tuned model to {repo_id}")
-
-        progress(1.0, desc="Complete!")
-        logs.append("\n🎉 Training complete!")
-
-    except Exception as e:
-        logs.append(f"\n❌ Error: {str(e)}")
-
-    return "\n".join(logs)
-
-
-# ------------------------------------------------------
-# 🎨 Gradio Interface
-# ------------------------------------------------------
+    tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=dataset.column_names)
+    output_log.append(f"   ✅ Tokenized {len(tokenized_dataset)} samples")
+
+    # ==== Training setup ====
+    progress(0.5, desc="Starting training...")
+    output_log.append("\n⚙️ Preparing Trainer...")
+
+    output_dir = "./Qwen-Gita-Checkpoints"
+    training_args = TrainingArguments(
+        output_dir=output_dir,
+        overwrite_output_dir=True,
+        per_device_train_batch_size=batch_size,
+        num_train_epochs=num_epochs,
+        learning_rate=learning_rate,
+        fp16=device == "cuda",
+        save_steps=100,
+        logging_steps=10,
+        save_total_limit=1,
+    )
+
+    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized_dataset,
+        tokenizer=tokenizer,
+        data_collator=data_collator,
+    )
+
+    # ==== Train ====
+    output_log.append("\n🚀 Training started ...")
+    trainer.train()
+    output_log.append("✅ Training complete!")
+
+    # ==== Push to Hugging Face Hub ====
+    repo_id = "rahul7star/Qwen0.5-3B-Gita"
+    output_log.append(f"\n☁️ Uploading to Hugging Face Hub: {repo_id}")
+    api = HfApi()
+    token = HfFolder.get_token()
+
+    model.push_to_hub(repo_id, token=token)
+    tokenizer.push_to_hub(repo_id, token=token)
+    output_log.append(f"✅ Model uploaded successfully to {repo_id}")
+
+    return "\n".join(output_log)
+
+# ---------------------------------------------------------------------
+# Gradio Interface
+# ---------------------------------------------------------------------
 def create_interface():
-    with gr.Blocks(title="PromptWizard — Qwen Gita Trainer") as demo:
+    with gr.Blocks(title="🧘 Qwen Gita Trainer") as demo:
         gr.Markdown("""
-        # 🧘 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`
+        # 🧘 Fine-tune Qwen 0.5B on Bhagavad Gita  
+        This app downloads `rahul7star/Gita`, trains the model to become a Gita teacher,  
+        and uploads results to `rahul7star/Qwen0.5-3B-Gita`.
         """)
 
-        with gr.Row():
-            with gr.Column():
-                gpu_status = gr.Textbox(
-                    label="GPU Status",
-                    value=check_gpu_status(),
-                    interactive=False,
-                )
-                model_name = gr.Textbox(
-                    value="Qwen/Qwen2.5-0.5B",
-                    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")
-                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 Logs",
-                    lines=25,
-                    max_lines=40,
-                    value="Click 'Start Fine-tuning' to train on Bhagavad Gita dataset...",
-                )
+        gpu_status = gr.Textbox(value=check_gpu_status(), label="GPU Status", interactive=False)
+        model_name = gr.Textbox(value="Qwen/Qwen2.5-0.5B", label="Base Model", 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")
+        learning_rate = gr.Number(value=5e-5, label="Learning Rate")
+        train_btn = gr.Button("🚀 Start Fine-tuning", variant="primary")
+        output = gr.Textbox(label="Training Log", lines=30)
 
         train_btn.click(
             fn=train_model,
@@ -238,20 +152,9 @@ def create_interface():
             outputs=output,
         )
 
-        gr.Markdown("""
-        ---
-        **Notes:**
-        - Downloads dataset: `rahul7star/Gita` → `/tmp/gita_data/Gita.csv`
-        - Trains using LoRA for efficiency
-        - Uploads to `rahul7star/Qwen0.5-3B-Gita`
-        """)
-
     return demo
 
+demo = create_interface()
 
-# ------------------------------------------------------
-# 🚪 Launch app
-# ------------------------------------------------------
 if __name__ == "__main__":
-    demo = create_interface()
     demo.launch()