Update app.py

app.py

@@ -1,7 +1,7 @@
 """
-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
+PromptWizard Qwen Training — Configurable Dataset & Repo
+Fine-tunes Qwen using a user-selected dataset and uploads the trained model
+to a user-specified Hugging Face Hub repo asynchronously with detailed logs.
 """
 
 import gradio as gr
@@ -13,16 +13,11 @@ from transformers import (
     Trainer,
     TrainingArguments,
 )
-from datasets import load_dataset, Dataset
+from datasets import load_dataset
 from peft import LoraConfig, get_peft_model, TaskType
 from huggingface_hub import HfApi, HfFolder, Repository
-import os, tempfile, shutil
-
-import asyncio
-import tempfile
-import shutil
-from huggingface_hub import HfApi, HfFolder, Repository
-import threading
+import os, tempfile, shutil, asyncio, threading, time
+from datetime import datetime
 
 # ==== Async upload wrapper ====
 def start_async_upload(local_dir, hf_repo, output_log):
@@ -30,7 +25,6 @@ def start_async_upload(local_dir, hf_repo, output_log):
         asyncio.run(async_upload_model(local_dir, hf_repo, output_log))
     threading.Thread(target=runner, daemon=True).start()
 
-
 async def async_upload_model(local_dir, hf_repo, output_log):
     try:
         token = HfFolder.get_token()
@@ -41,7 +35,6 @@ async def async_upload_model(local_dir, hf_repo, output_log):
 
         with tempfile.TemporaryDirectory() as tmpdir:
             repo = Repository(local_dir=tmpdir, clone_from=hf_repo, use_auth_token=token)
-            # Copy model files
             shutil.copytree(local_dir, tmpdir, dirs_exist_ok=True)
             repo.push_to_hub(commit_message="Upload fine-tuned model")
 
@@ -49,40 +42,38 @@ async def async_upload_model(local_dir, hf_repo, output_log):
     except Exception as e:
         output_log.append(f"\n❌ Async upload error: {e}")
 
-
-# === GPU check (Zero GPU compatible) ===
+# ==== GPU check ====
 def check_gpu_status():
     return "🚀 Zero GPU Ready - GPU will be allocated when training starts"
 
+# ==== Logging helper ====
+def log_message(output_log, msg):
+    line = f"[{datetime.now().strftime('%H:%M:%S')}] {msg}"
+    print(line)
+    output_log.append(line)
 
-# === Main Training ===
+# ==== Main Training ====
 @spaces.GPU(duration=300)
-def train_model(model_name, num_epochs, batch_size, learning_rate, progress=gr.Progress()):
-    progress(0, desc="Initializing...")
+def train_model(base_model, dataset_name, num_epochs, batch_size, learning_rate, hf_repo):
     output_log = []
-
     try:
-        # ==== Device ====
+        log_message(output_log, "🔍 Initializing training sequence...")
+
+        # ===== Device =====
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        output_log.append(f"🎮 Using device: {device}")
+        log_message(output_log, f"🎮 Using device: {device}")
         if device == "cuda":
-            output_log.append(f"✅ GPU: {torch.cuda.get_device_name(0)}")
+            log_message(output_log, 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...")
+        # ===== Load dataset =====
+        log_message(output_log, f"\n📚 Loading dataset: {dataset_name} ...")
+        dataset = load_dataset(dataset_name, split="train")
+        log_message(output_log, f"   Loaded {len(dataset)} samples")
+        log_message(output_log, f"   Columns: {dataset.column_names}")
 
-        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 ====
+        # ===== Format examples =====
         def format_example(item):
-            text = (
-                item.get("text")
-                or item.get("content")
-                or " ".join(str(v) for v in item.values())
-            )
+            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|>
@@ -90,33 +81,27 @@ You are a wise teacher interpreting Bhagavad Gita with deep insights.
 <|assistant|>
 """
             return {"text": prompt}
-
         dataset = dataset.map(format_example)
-        output_log.append(f"   ✅ Formatted {len(dataset)} examples")
+        log_message(output_log, f"✅ Formatted {len(dataset)} examples")
 
-        # ==== Model & Tokenizer ====
-        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)
+        # ===== Load model & tokenizer =====
+        log_message(output_log, f"\n🤖 Loading model: {base_model}")
+        tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
-
         model = AutoModelForCausalLM.from_pretrained(
-            model_name,
+            base_model,
             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...")
+        log_message(output_log, "✅ Model and tokenizer loaded successfully")
+        log_message(output_log, f"Tokenizer vocab size: {tokenizer.vocab_size}")
 
+        # ===== LoRA configuration =====
+        log_message(output_log, "\n⚙️ Configuring LoRA for efficient fine-tuning...")
         lora_config = LoraConfig(
             task_type=TaskType.CAUSAL_LM,
             r=8,
@@ -126,12 +111,10 @@ You are a wise teacher interpreting Bhagavad Gita with deep insights.
             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:,}")
+        log_message(output_log, f"Trainable params after LoRA: {trainable_params:,}")
 
-        # ==== Tokenization + Labels ====
-        progress(0.5, desc="Tokenizing dataset...")
+        # ===== Tokenization + labels =====
         def tokenize_fn(examples):
             tokenized = tokenizer(
                 examples["text"],
@@ -139,15 +122,12 @@ You are a wise teacher interpreting Bhagavad Gita with deep insights.
                 truncation=True,
                 max_length=256,
             )
-            # Add labels for causal LM
             tokenized["labels"] = tokenized["input_ids"].copy()
             return tokenized
-
         dataset = dataset.map(tokenize_fn, batched=True)
-        output_log.append("   ✅ Tokenization + labels done")
+        log_message(output_log, "✅ Tokenization + labels done")
 
-        # ==== Training arguments ====
-        progress(0.6, desc="Setting up training...")
+        # ===== Training arguments =====
         output_dir = "./qwen-gita-lora"
         training_args = TrainingArguments(
             output_dir=output_dir,
@@ -162,7 +142,6 @@ You are a wise teacher interpreting Bhagavad Gita with deep insights.
             learning_rate=learning_rate,
             max_steps=100,
         )
-
         trainer = Trainer(
             model=model,
             args=training_args,
@@ -170,54 +149,38 @@ You are a wise teacher interpreting Bhagavad Gita with deep insights.
             tokenizer=tokenizer,
         )
 
-       # ==== Train ====
-        progress(0.7, desc="Training...")
-        output_log.append("\n🚀 Starting training...\n" + "=" * 50)
-        train_result = trainer.train()
-
-# ==== Save model locally ====
-        progress(0.85, desc="Saving model...")
-        output_log.append("\n💾 Saving model locally...")
+        # ===== Train =====
+        log_message(output_log, "\n🚀 Starting training...")
+        trainer.train()
+        log_message(output_log, "\n💾 Saving trained model locally...")
         trainer.save_model(output_dir)
         tokenizer.save_pretrained(output_dir)
 
-# ==== Async upload ====
-        # ==== Async upload ====
-        hf_repo = "rahul7star/Qwen0.5-3B-Gita"
+        # ===== Async upload to repo from UI input =====
+        log_message(output_log, f"\n☁️ Initiating async upload to {hf_repo}")
         start_async_upload(output_dir, hf_repo, output_log)
 
-        progress(1.0, desc="Complete!")
-       
-
-        output_log.append("\n✅ Training complete & model uploaded successfully!")
+        log_message(output_log, "✅ Training complete & async upload started!")
 
     except Exception as e:
-        output_log.append(f"\n❌ Error: {e}")
+        log_message(output_log, f"\n❌ Error during training: {e}")
 
     return "\n".join(output_log)
 
-
-# === Gradio Interface ===
+# ==== Gradio Interface ====
 def create_interface():
-    with gr.Blocks(title="PromptWizard — Qwen Gita Trainer") as demo:
+    with gr.Blocks(title="PromptWizard — Qwen Trainer") as demo:
         gr.Markdown("""
-        # 🧘 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**.
+        # 🧘 PromptWizard Qwen Fine-tuning  
+        Fine-tune Qwen on any dataset and upload to any Hugging Face repo.
         """)
 
         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,
-                )
+                gr.Textbox(label="GPU Status", value=check_gpu_status(), interactive=False)
+                base_model = gr.Textbox(label="Base Model", value="Qwen/Qwen2.5-0.5B")
+                dataset_name = gr.Textbox(label="Dataset Name", value="rahul7star/Gita")
+                hf_repo = gr.Textbox(label="HF Repo for Upload", value="rahul7star/Qwen0.5-3B-Gita")
                 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")
@@ -228,18 +191,17 @@ def create_interface():
                     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.",
+                    value="Click 'Start Fine-tuning' to train and upload your model.",
                 )
 
         train_btn.click(
             fn=train_model,
-            inputs=[model_name, num_epochs, batch_size, learning_rate],
+            inputs=[base_model, dataset_name, num_epochs, batch_size, learning_rate, hf_repo],
             outputs=output,
         )
 
     return demo
 
-
 if __name__ == "__main__":
     demo = create_interface()
-    demo.launch()
+    demo.launch(server_name="0.0.0.0", server_port=7860)