Update app.py

app.py

@@ -20,13 +20,8 @@ import os, tempfile, shutil, asyncio, threading, time
 from datetime import datetime
 
 # ==== Async upload wrapper ====
-import threading
-import asyncio
-
 def start_async_upload(local_dir, hf_repo, output_log):
-    """
-    Starts async model upload in a background thread.
-    """
+    """Starts async model upload in a background thread."""
     def runner():
         output_log.append(f"[INFO] 🚀 Async upload thread started for repo: {hf_repo}")
         asyncio.run(async_upload_model(local_dir, hf_repo, output_log))
@@ -35,16 +30,8 @@ def start_async_upload(local_dir, hf_repo, output_log):
     threading.Thread(target=runner, daemon=True).start()
 
 
-import os
-import shutil
-import tempfile
-import asyncio
-from huggingface_hub import HfApi, HfFolder, Repository
-
 async def async_upload_model(local_dir, hf_repo, output_log, max_retries=3):
-    """
-    Uploads a local model directory to HF Hub asynchronously with detailed logging.
-    """
+    """Uploads a local model directory to HF Hub asynchronously with detailed logging."""
     try:
         token = HfFolder.get_token()
         api = HfApi()
@@ -97,6 +84,7 @@ async def async_upload_model(local_dir, hf_repo, output_log, max_retries=3):
     except Exception as e:
         output_log.append(f"[ERROR] ❌ Unexpected error during upload: {e}")
 
+
 # ==== GPU check ====
 def check_gpu_status():
     return "🚀 Zero GPU Ready - GPU will be allocated when training starts"
@@ -107,18 +95,12 @@ def log_message(output_log, msg):
     print(line)
     output_log.append(line)
 
+
 # ==== Main Training ====
 @spaces.GPU(duration=300)
 def train_model(base_model, dataset_name, num_epochs, batch_size, learning_rate, hf_repo):
-    """
-    Fine-tune a base model using LoRA with train/test split and async upload.
-    """
     output_log = []
-    test_split=0.2
-
-    def log_message(log_list, msg):
-        print(msg)
-        log_list.append(msg)
+    test_split = 0.2
 
     try:
         log_message(output_log, "🔍 Initializing training sequence...")
@@ -132,7 +114,6 @@ def train_model(base_model, dataset_name, num_epochs, batch_size, learning_rate,
         # ===== Load dataset =====
         log_message(output_log, f"\n📚 Loading dataset: {dataset_name} ...")
         dataset = load_dataset(dataset_name)
-        # Determine train/test split
         dataset = dataset["train"].train_test_split(test_size=test_split)
         train_dataset = dataset["train"]
         test_dataset = dataset["test"]
@@ -212,11 +193,10 @@ You are a wise teacher interpreting Bhagavad Gita with deep insights.
             warmup_steps=10,
             logging_steps=5,
             save_strategy="epoch",
-            evaluation_strategy="epoch",
             fp16=device == "cuda",
             optim="adamw_torch",
             learning_rate=learning_rate,
-            max_steps=100,
+            max_steps=100,  # Limit for demo
         )
 
         trainer = Trainer(
@@ -245,6 +225,7 @@ You are a wise teacher interpreting Bhagavad Gita with deep insights.
 
     return "\n".join(output_log)
 
+
 # ==== Gradio Interface ====
 def create_interface():
     with gr.Blocks(title="PromptWizard — Qwen Trainer") as demo:
@@ -280,6 +261,7 @@ def create_interface():
 
     return demo
 
+
 if __name__ == "__main__":
     demo = create_interface()
     demo.launch(server_name="0.0.0.0", server_port=7860)