Update app.py

app.py

@@ -1,183 +1,213 @@
+"""
+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 os
 import torch
-import gradio as gr
-from datasets import load_dataset
 from transformers import (
-    AutoTokenizer,
     AutoModelForCausalLM,
+    AutoTokenizer,
     Trainer,
     TrainingArguments,
-    DataCollatorForLanguageModeling,
 )
-from huggingface_hub import HfApi, HfFolder
+from datasets import load_dataset, Dataset
+from peft import LoraConfig, get_peft_model, TaskType
+from huggingface_hub import HfApi, HfFolder, Repository
+import os, tempfile, shutil
 
-# ---------------------------------------------------------------------
-# GPU check
-# ---------------------------------------------------------------------
-device = "cuda" if torch.cuda.is_available() else "cpu"
 
+# === GPU check (Zero GPU compatible) ===
 def check_gpu_status():
-    return f"✅ GPU: {torch.cuda.get_device_name(0)}" if device == "cuda" else "⚠️ Using CPU only"
-# ------------------------------------------------------
-# 🧩 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
-# ------------------------------------------------------
-
-# ---------------------------------------------------------------------
-# Training Logic
-# ---------------------------------------------------------------------
+    return "🚀 Zero GPU Ready - GPU will be allocated when training starts"
+
+
+# === Main Training ===
 @spaces.GPU(duration=300)
-def train_model(model_name, num_epochs, batch_size, learning_rate, progress=gr.Progress(track_tqdm=True)):
+def train_model(model_name, num_epochs, batch_size, learning_rate, progress=gr.Progress()):
+    progress(0, desc="Initializing...")
     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())
-        )
-        prompt = f"""<|system|>
+    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 from CSV")
+        output_log.append(f"   Columns: {dataset.column_names}")
+
+        # ==== Format data ====
+        def format_example(item):
+            # 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 wise teacher interpreting Bhagavad Gita with deep insights.
 <|user|>
 {text}
 <|assistant|>
 """
-        return {"text": prompt}
+            return {"text": prompt}
 
-    dataset = dataset.map(format_example)
-    output_log.append(f"   ✅ Formatted {len(dataset)} examples")
+        dataset = dataset.map(format_example)
+        output_log.append(f"   ✅ Formatted {len(dataset)} examples")
 
-    # ==== Load tokenizer & model ====
-    progress(0.3, desc="Loading model and tokenizer...")
-    output_log.append("\n🤖 Loading Qwen model and tokenizer...")
+        # ==== Model ====
+        progress(0.3, desc="Loading model & tokenizer...")
+        model_name = "Qwen/Qwen2.5-0.5B"
+        output_log.append(f"\n���� Loading model: {model_name}")
 
-    base_model = "Qwen/Qwen2.5-0.5B"
-    tokenizer = AutoTokenizer.from_pretrained(base_model)
+        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
 
-    # Fix missing pad token
-    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 ====
+        progress(0.4, desc="Configuring LoRA...")
+        output_log.append("\n⚙️ Setting up LoRA for efficient fine-tuning...")
+
+        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:,}")
+
+        # ==== Tokenization ====
+        progress(0.5, desc="Tokenizing dataset...")
+        def tokenize_fn(examples):
+            return tokenizer(
+                examples["text"],
+                padding="max_length",
+                truncation=True,
+                max_length=256,
+            )
+
+        dataset = dataset.map(tokenize_fn, batched=True)
+        output_log.append("   ✅ Tokenization done")
+
+        # ==== Training arguments ====
+        progress(0.6, desc="Setting up training...")
+        output_dir = "./qwen-gita-lora"
+        training_args = TrainingArguments(
+            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="epoch",
+            fp16=device == "cuda",
+            optim="adamw_torch",
+            learning_rate=learning_rate,
+            max_steps=100,
+        )
 
-    model = AutoModelForCausalLM.from_pretrained(
-        base_model,
-        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-    ).to(device)
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=dataset,
+            tokenizer=tokenizer,
+        )
 
-    # ==== Tokenize dataset ====
-    progress(0.4, desc="Tokenizing dataset...")
-    output_log.append("\n✏️ Tokenizing dataset...")
+        # ==== Train ====
+        progress(0.7, desc="Training...")
+        output_log.append("\n🚀 Starting training...\n" + "=" * 50)
+        train_result = trainer.train()
 
-    def tokenize_function(examples):
-        return tokenizer(
-            examples["text"],
-            truncation=True,
-            padding="max_length",
-            max_length=512,
-        )
+        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}")
 
-    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}")
+        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 Interface
-# ---------------------------------------------------------------------
+
+# === Gradio Interface ===
 def create_interface():
-    with gr.Blocks(title="🧘 Qwen Gita Trainer") as demo:
+    with gr.Blocks(title="PromptWizard — Qwen Gita Trainer") as demo:
         gr.Markdown("""
-        # 🧘 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`.
+        # 🧘 PromptWizard Qwen Fine-tuning — Gita Edition  
+        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**.
         """)
 
-        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)
+        with gr.Row():
+            with gr.Column():
+                gpu_status = gr.Textbox(
+                    label="GPU Status",
+                    value=check_gpu_status(),
+                    interactive=False,
+                )
+                model_name = gr.Textbox(
+                    label="Base Model",
+                    value="Qwen/Qwen2.5-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 Fine-tuning", variant="primary")
+
+            with gr.Column():
+                output = gr.Textbox(
+                    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,
@@ -187,7 +217,7 @@ def create_interface():
 
     return demo
 
-demo = create_interface()
 
 if __name__ == "__main__":
+    demo = create_interface()
     demo.launch()