Update app.py

app.py

@@ -1,70 +1,74 @@
 """
-PromptWizard Qwen Training using rahul7star/Gita dataset (.csv)
-Optimized for HuggingFace Spaces with Zero GPU allocation
+PromptWizard Qwen Training — Gita Edition
+Fine-tunes Qwen using rahul7star/Gita dataset (.csv)
+Uploads trained model to rahul7star/Qwen0.5-3B-Gita on Hugging Face Hub
 """
 
 import gradio as gr
 import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    Trainer,
+    TrainingArguments,
+)
 from datasets import load_dataset, Dataset
 from peft import LoraConfig, get_peft_model, TaskType
-import os
+from huggingface_hub import HfApi, HfFolder, Repository
+import os, tempfile, shutil
 
-# GPU check (Zero GPU compatible)
+
+# === 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
+
+# === Main Training ===
+@spaces.GPU(duration=300)
 def train_model(model_name, num_epochs, batch_size, learning_rate, progress=gr.Progress()):
     progress(0, desc="Initializing...")
     output_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 your dataset ===
-        progress(0.1, desc="Loading Gita dataset...")
-        output_log.append("\n📚 Loading rahul7star/Gita dataset from Hugging Face Hub...")
+        # ==== Load dataset ====
+        progress(0.1, desc="Loading rahul7star/Gita dataset...")
+        output_log.append("\n📚 Loading dataset from rahul7star/Gita...")
 
-        # 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"   Loaded {len(dataset)} samples from CSV")
         output_log.append(f"   Columns: {dataset.column_names}")
 
-        # === Format each row into text prompt ===
+        # ==== Format data ====
         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())
+            # Use "text" or "content" column if available
+            text = (
+                item.get("text")
+                or item.get("content")
+                or " ".join(str(v) for v in item.values())
+            )
 
             prompt = f"""<|system|>
-You are a spiritual teacher interpreting Bhagavad Gita verses with deep meaning.
+You are a wise teacher interpreting Bhagavad Gita with deep insights.
 <|user|>
-{ text }
+{text}
 <|assistant|>
 """
             return {"text": prompt}
 
-        train_dataset = dataset.map(format_example)
-        output_log.append(f"   Formatted {len(train_dataset)} training samples")
+        dataset = dataset.map(format_example)
+        output_log.append(f"   ✅ Formatted {len(dataset)} examples")
 
-        # === 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}...")
+        # ==== Model ====
+        progress(0.3, desc="Loading model & tokenizer...")
+        model_name = "Qwen/Qwen2.5-0.5B"
+        output_log.append(f"\n🤖 Loading model: {model_name}")
 
         tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
         if tokenizer.pad_token is None:
@@ -74,16 +78,16 @@ You are a spiritual teacher interpreting Bhagavad Gita verses with deep meaning.
             model_name,
             trust_remote_code=True,
             torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-            low_cpu_mem_usage=True
+            low_cpu_mem_usage=True,
         )
-
         if device == "cuda":
             model = model.to(device)
+
         output_log.append("   ✅ Model loaded successfully")
 
-        # === LoRA configuration ===
+        # ==== LoRA ====
         progress(0.4, desc="Configuring LoRA...")
-        output_log.append("\n⚙️ Setting up LoRA for efficient training...")
+        output_log.append("\n⚙️ Setting up LoRA for efficient fine-tuning...")
 
         lora_config = LoraConfig(
             task_type=TaskType.CAUSAL_LM,
@@ -91,72 +95,93 @@ You are a spiritual teacher interpreting Bhagavad Gita verses with deep meaning.
             lora_alpha=16,
             lora_dropout=0.1,
             target_modules=["q_proj", "v_proj"],
-            bias="none"
+            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...")
+        # ==== Tokenization ====
+        progress(0.5, desc="Tokenizing dataset...")
         def tokenize_fn(examples):
             return tokenizer(
                 examples["text"],
                 padding="max_length",
                 truncation=True,
-                max_length=256
+                max_length=256,
             )
-        train_dataset = train_dataset.map(tokenize_fn, batched=True)
+
+        dataset = dataset.map(tokenize_fn, batched=True)
         output_log.append("   ✅ Tokenization done")
 
-        # === Training configuration ===
+        # ==== Training arguments ====
         progress(0.6, desc="Setting up training...")
+        output_dir = "./qwen-gita-lora"
         training_args = TrainingArguments(
-            output_dir="./qwen-gita-lora",
+            output_dir=output_dir,
             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",
+            save_strategy="epoch",
             fp16=device == "cuda",
             optim="adamw_torch",
             learning_rate=learning_rate,
-            max_steps=50,  # Short demo
+            max_steps=100,
         )
 
         trainer = Trainer(
             model=model,
             args=training_args,
-            train_dataset=train_dataset,
+            train_dataset=dataset,
             tokenizer=tokenizer,
         )
 
-        # === Train ===
-        progress(0.7, desc="Training model...")
-        output_log.append("\n🚀 Starting training...\n" + "="*50)
+        # ==== Train ====
+        progress(0.7, desc="Training...")
+        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(0.85, desc="Saving model...")
+        output_log.append("\n💾 Saving model locally...")
+        trainer.save_model(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}")
+
+        api = HfApi()
+        token = HfFolder.get_token()
+
+        # Create repo if not exists
+        api.create_repo(repo_id=hf_repo, exist_ok=True)
+
+        # Clone & push
+        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 LoRA model")
 
         progress(1.0, desc="Complete!")
+        output_log.append("\n✅ Training complete & model uploaded successfully!")
+
     except Exception as e:
         output_log.append(f"\n❌ Error: {e}")
 
     return "\n".join(output_log)
 
 
-# === Gradio UI ===
+# === Gradio Interface ===
 def create_interface():
-    with gr.Blocks(title="PromptWizard Gita Trainer") as demo:
+    with gr.Blocks(title="PromptWizard — Qwen 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.  
+        Fine-tune **Qwen 0.5B** on your dataset [rahul7star/Gita](https://huggingface.co/datasets/rahul7star/Gita)  
+        and auto-upload to your model repo **rahul7star/Qwen0.5-3B-Gita**.
         """)
 
         with gr.Row():
@@ -164,32 +189,30 @@ def create_interface():
                 gpu_status = gr.Textbox(
                     label="GPU Status",
                     value=check_gpu_status(),
-                    interactive=False
+                    interactive=False,
                 )
-
-                model_name = gr.Dropdown(
-                    choices=["Qwen/Qwen2.5-0.5B", "Qwen/Qwen2.5-1.5B"],
+                model_name = gr.Textbox(
+                    label="Base Model",
                     value="Qwen/Qwen2.5-0.5B",
-                    label="Model (recommended: 0.5B)"
+                    interactive=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 Training", variant="primary")
+                train_btn = gr.Button("🚀 Start Fine-tuning", 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."
+                    label="Training Log",
+                    lines=25,
+                    max_lines=40,
+                    value="Click 'Start Fine-tuning' to train on the Gita dataset and upload to your model repo.",
                 )
 
         train_btn.click(
             fn=train_model,
             inputs=[model_name, num_epochs, batch_size, learning_rate],
-            outputs=output
+            outputs=output,
         )
 
     return demo
@@ -197,4 +220,4 @@ def create_interface():
 
 if __name__ == "__main__":
     demo = create_interface()
-    demo.launch()
+    demo.launch()