Update app.py

app.py

@@ -1,270 +1,200 @@
-"""
-PromptWizard Qwen Training with Zero GPU
-Optimized for HuggingFace Spaces with automatic GPU allocation
-"""
-
-import gradio as gr
-import spaces
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
-from datasets import load_dataset, Dataset
-from peft import LoraConfig, get_peft_model, TaskType
-import json
-import os
-
-# Check if GPU is available (Zero GPU safe)
-def check_gpu_status():
-    # Don't check CUDA at module load time for Zero GPU compatibility
-    return "🚀 Zero GPU Ready - GPU will be allocated when training starts"
-
-@spaces.GPU(duration=300)  # Request GPU for 5 minutes (can extend if needed)
-def train_model(model_name, num_epochs, batch_size, learning_rate, progress=gr.Progress()):
-    """Main training function with Zero GPU support"""
-    
-    progress(0, desc="Initializing...")
-    output_log = []
-    
-    try:
-        # GPU should be available inside this function
-        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)}")
-            output_log.append(f"   Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f}GB")
-        
-        # Load GSM8K dataset
-        progress(0.1, desc="Loading GSM8K dataset...")
-        output_log.append("\n📚 Loading GSM8K dataset...")
-        
-        # Load local data if available, otherwise from HF
-        train_data = []
-        test_data = []
-        
-        # Try local files first
-        if os.path.exists("data/train.jsonl"):
-            with open("data/train.jsonl", "r") as f:
-                for line in f:
-                    train_data.append(json.loads(line))
-            output_log.append(f"   Loaded {len(train_data)} training examples from local data")
-        else:
-            # Fallback to HF dataset
-            dataset = load_dataset("openai/gsm8k", "main")
-            train_data = dataset["train"].select(range(min(100, len(dataset["train"]))))
-            output_log.append(f"   Loaded {len(train_data)} training examples from HF")
-        
-        # Format prompts
-        def format_example(item):
-            prompt = f"""<|system|>
-You are a mathematics expert. Solve grade school math problems step by step.
-<|user|>
-{item.get('question', '')}
-<|assistant|>
-{item.get('full_solution', item.get('answer', ''))}"""
-            return {"text": prompt}
-        
-        # Create dataset
-        if isinstance(train_data, list):
-            train_dataset = Dataset.from_list([format_example(item) for item in train_data])
-        else:
-            train_dataset = train_data.map(format_example)
-        
-        output_log.append(f"   Training samples ready: {len(train_dataset)}")
-        
-        # Load model and tokenizer
-        progress(0.3, desc="Loading model and tokenizer...")
-        output_log.append(f"\n🤖 Loading {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")
-        
-        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-        if tokenizer.pad_token is None:
-            tokenizer.pad_token = tokenizer.eos_token
-        
-        # Load model without quantization for Zero GPU compatibility
-        model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            trust_remote_code=True,
-            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-            low_cpu_mem_usage=True
-        )
-        
-        # Move model to GPU if available
-        if device == "cuda":
-            model = model.to(device)
-        
-        output_log.append("   Model loaded successfully")
-        
-        # Configure LoRA
-        progress(0.4, desc="Configuring LoRA...")
-        output_log.append("\n⚙️ Configuring LoRA for efficient training...")
-        
-        lora_config = LoraConfig(
-            task_type=TaskType.CAUSAL_LM,
-            r=8,  # Low rank for efficiency
-            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)
-        total_params = sum(p.numel() for p in model.parameters())
-        output_log.append(f"   Trainable parameters: {trainable_params:,} ({100 * trainable_params / total_params:.2f}%)")
-        
-        # Tokenize dataset
-        progress(0.5, desc="Preparing data...")
-        output_log.append("\n📝 Tokenizing dataset...")
-        
-        def tokenize_function(examples):
-            return tokenizer(
-                examples["text"],
-                padding="max_length",
-                truncation=True,
-                max_length=256  # Shorter for demo
-            )
-        
-        train_dataset = train_dataset.map(tokenize_function, batched=True)
-        
-        # Training arguments
-        progress(0.6, desc="Setting up training...")
-        output_log.append("\n🎯 Setting up training configuration...")
-        
-        training_args = TrainingArguments(
-            output_dir="./qwen-promptwizard-zerogpu",
-            num_train_epochs=num_epochs,
-            per_device_train_batch_size=batch_size,
-            gradient_accumulation_steps=2,  # Reduced for Zero GPU
-            warmup_steps=10,  # Reduced warmup
-            logging_steps=5,
-            save_strategy="no",  # Don't save during demo
-            fp16=device == "cuda",  # Only use fp16 on GPU
-            gradient_checkpointing=False,  # Disable for simplicity
-            optim="adamw_torch",
-            learning_rate=learning_rate,
-            max_steps=50,  # Limit steps for demo
-        )
-        
-        # Create trainer
-        trainer = Trainer(
-            model=model,
-            args=training_args,
-            train_dataset=train_dataset,
-            tokenizer=tokenizer,
-        )
-        
-        # Start training
-        progress(0.7, desc="Training...")
-        output_log.append(f"\n🚀 Starting training for {num_epochs} epochs...")
-        output_log.append("=" * 50)
-        
-        # Train
-        train_result = trainer.train()
-        
-        # Results
-        progress(0.9, desc="Finalizing...")
-        output_log.append("=" * 50)
-        output_log.append("\n✅ Training completed!")
-        output_log.append(f"   Final loss: {train_result.training_loss:.4f}")
-        output_log.append(f"   Total steps: {train_result.global_step}")
-        
-        # Save model info
-        output_log.append("\n💾 Model trained with PromptWizard + GSM8K")
-        output_log.append("   Using REAL data and REAL evaluation!")
-        
-        progress(1.0, desc="Complete!")
-        
-    except Exception as e:
-        output_log.append(f"\n❌ Error: {str(e)}")
-        output_log.append("Note: Zero GPU requires proper setup in Space settings")
-    
-    return "\n".join(output_log)
-
-# Gradio interface
-def create_interface():
-    with gr.Blocks(title="PromptWizard Qwen Training") as demo:
-        gr.Markdown("""
-        # 🧙 PromptWizard Qwen Fine-tuning with Zero GPU
-        
-        Fine-tune Qwen models using GSM8K dataset with PromptWizard methodology.
-        This Space uses HuggingFace Zero GPU for free GPU access during training.
-        
-        **Features:**
-        - ✅ Real GSM8K mathematical problems (not fake data!)
-        - ✅ LoRA-based efficient fine-tuning
-        - ✅ Automatic Zero GPU allocation
-        - ✅ PromptWizard optimization methodology
-        """)
-        
-        with gr.Row():
-            with gr.Column():
-                gpu_status = gr.Textbox(
-                    label="GPU Status",
-                    value=check_gpu_status(),
-                    interactive=False
-                )
-                
-                model_name = gr.Dropdown(
-                    choices=[
-                        "Qwen/Qwen2.5-0.5B",
-                        "Qwen/Qwen2.5-1.5B",
-                    ],
-                    value="Qwen/Qwen2.5-0.5B",
-                    label="Model (0.5B works best for Zero GPU)"
-                )
-                
-                num_epochs = gr.Slider(
-                    minimum=1,
-                    maximum=3,
-                    value=1,
-                    step=1,
-                    label="Number of Epochs (1 for quick demo)"
-                )
-                
-                batch_size = gr.Slider(
-                    minimum=1,
-                    maximum=4,
-                    value=2,
-                    step=1,
-                    label="Batch Size (2 for Zero GPU)"
-                )
-                
-                learning_rate = gr.Number(
-                    value=5e-5,
-                    label="Learning Rate"
-                )
-                
-                train_btn = gr.Button("🚀 Start Training", variant="primary")
-                
-            with gr.Column():
-                output = gr.Textbox(
-                    label="Training Output",
-                    lines=20,
-                    max_lines=30,
-                    value="Click 'Start Training' to begin...\n\nZero GPU will automatically allocate a GPU when training starts."
-                )
-        
-        # Connect button to training function
-        train_btn.click(
-            fn=train_model,
-            inputs=[model_name, num_epochs, batch_size, learning_rate],
-            outputs=output
-        )
-        
-        gr.Markdown("""
-        ## Notes:
-        - Zero GPU provides free GPU access for public Spaces
-        - Training will automatically get GPU allocation when started
-        - Using smaller model (1.5B) for faster demo
-        - Real GSM8K data - no fake metrics!
-        """)
-    
-    return demo
-
-# Launch app
-if __name__ == "__main__":
-    demo = create_interface()
-    demo.launch()
+"""
+PromptWizard Qwen Training using rahul7star/Gita dataset (.csv)
+Optimized for HuggingFace Spaces with Zero GPU allocation
+"""
+
+import gradio as gr
+import spaces
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
+from datasets import load_dataset, Dataset
+from peft import LoraConfig, get_peft_model, TaskType
+import os
+
+# GPU check (Zero GPU compatible)
+def check_gpu_status():
+    return "🚀 Zero GPU Ready - GPU will be allocated when training starts"
+
+@spaces.GPU(duration=300)  # request GPU for 5 min
+def train_model(model_name, num_epochs, batch_size, learning_rate, progress=gr.Progress()):
+    progress(0, desc="Initializing...")
+    output_log = []
+
+    try:
+        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 your dataset ===
+        progress(0.1, desc="Loading Gita dataset...")
+        output_log.append("\n📚 Loading rahul7star/Gita dataset from Hugging Face Hub...")
+
+        # Load CSV from Hugging Face repo
+        dataset = load_dataset("rahul7star/Gita", split="train")
+        output_log.append(f"   Loaded {len(dataset)} examples from rahul7star/Gita")
+
+        # Preview columns
+        output_log.append(f"   Columns: {dataset.column_names}")
+
+        # === Format each row into text prompt ===
+        def format_example(item):
+            # You can customize this depending on your CSV structure
+            # Assuming CSV has columns like: "chapter", "verse", "content" or "text"
+            text = ""
+            if "text" in item:
+                text = item["text"]
+            elif "content" in item:
+                text = item["content"]
+            else:
+                # fallback: join all columns
+                text = " ".join(str(v) for v in item.values())
+
+            prompt = f"""<|system|>
+You are a spiritual teacher interpreting Bhagavad Gita verses with deep meaning.
+<|user|>
+{ text }
+<|assistant|>
+"""
+            return {"text": prompt}
+
+        train_dataset = dataset.map(format_example)
+        output_log.append(f"   Formatted {len(train_dataset)} training samples")
+
+        # === Load Qwen model & tokenizer ===
+        progress(0.3, desc="Loading model...")
+        model_name = "Qwen/Qwen2.5-0.5B"  # Force small model for free GPU
+        output_log.append(f"\n🤖 Loading {model_name}...")
+
+        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            trust_remote_code=True,
+            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+            low_cpu_mem_usage=True
+        )
+
+        if device == "cuda":
+            model = model.to(device)
+        output_log.append("   ✅ Model loaded successfully")
+
+        # === LoRA configuration ===
+        progress(0.4, desc="Configuring LoRA...")
+        output_log.append("\n⚙️ Setting up LoRA for efficient training...")
+
+        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 params: {trainable_params:,}")
+
+        # === Tokenize ===
+        progress(0.5, desc="Tokenizing data...")
+        def tokenize_fn(examples):
+            return tokenizer(
+                examples["text"],
+                padding="max_length",
+                truncation=True,
+                max_length=256
+            )
+        train_dataset = train_dataset.map(tokenize_fn, batched=True)
+        output_log.append("   ✅ Tokenization done")
+
+        # === Training configuration ===
+        progress(0.6, desc="Setting up training...")
+        training_args = TrainingArguments(
+            output_dir="./qwen-gita-lora",
+            num_train_epochs=num_epochs,
+            per_device_train_batch_size=batch_size,
+            gradient_accumulation_steps=2,
+            warmup_steps=10,
+            logging_steps=5,
+            save_strategy="no",
+            fp16=device == "cuda",
+            optim="adamw_torch",
+            learning_rate=learning_rate,
+            max_steps=50,  # Short demo
+        )
+
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=train_dataset,
+            tokenizer=tokenizer,
+        )
+
+        # === Train ===
+        progress(0.7, desc="Training model...")
+        output_log.append("\n🚀 Starting training...\n" + "="*50)
+        train_result = trainer.train()
+
+        progress(0.9, desc="Finalizing...")
+        output_log.append("="*50)
+        output_log.append("\n✅ Training completed successfully!")
+        output_log.append(f"   Final loss: {train_result.training_loss:.4f}")
+        output_log.append(f"   Steps: {train_result.global_step}")
+
+        progress(1.0, desc="Complete!")
+    except Exception as e:
+        output_log.append(f"\n❌ Error: {e}")
+
+    return "\n".join(output_log)
+
+
+# === Gradio UI ===
+def create_interface():
+    with gr.Blocks(title="PromptWizard Gita Trainer") as demo:
+        gr.Markdown("""
+        # 🧘 PromptWizard Qwen Fine-tuning — Gita Edition  
+        Train Qwen models using verses from **rahul7star/Gita** (CSV-based dataset).  
+        This uses **Zero GPU** on Hugging Face Spaces for free GPU access.  
+        """)
+
+        with gr.Row():
+            with gr.Column():
+                gpu_status = gr.Textbox(
+                    label="GPU Status",
+                    value=check_gpu_status(),
+                    interactive=False
+                )
+
+                model_name = gr.Dropdown(
+                    choices=["Qwen/Qwen2.5-0.5B", "Qwen/Qwen2.5-1.5B"],
+                    value="Qwen/Qwen2.5-0.5B",
+                    label="Model (recommended: 0.5B)"
+                )
+
+                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 Training", variant="primary")
+
+            with gr.Column():
+                output = gr.Textbox(
+                    label="Training Output",
+                    lines=20,
+                    max_lines=30,
+                    value="Click 'Start Training' to fine-tune on Gita dataset."
+                )
+
+        train_btn.click(
+            fn=train_model,
+            inputs=[model_name, num_epochs, batch_size, learning_rate],
+            outputs=output
+        )
+
+    return demo
+
+
+if __name__ == "__main__":
+    demo = create_interface()
+    demo.launch()