Update app.py

app.py

@@ -1,41 +1,59 @@
 import gradio as gr
 from huggingface_hub import hf_hub_download
 from llama_cpp import Llama
+import threading
+import time
 
-# Global variable to hold the model (starts empty)
+# --- Global State ---
 llm_model = None
+load_status = "Starting..."
+is_loaded = False
 
-def load_model():
-    """Loads the model only when needed."""
-    global llm_model
+# --- Background Loader ---
+def load_model_in_background():
+    global llm_model, load_status, is_loaded
     
-    # If already loaded, just return it
-    if llm_model is not None:
-        return llm_model
+    try:
+        print("⏳ Background thread started...")
+        load_status = "⬇️ Downloading model (approx 1-2 mins)..."
+        
+        # Download the model
+        model_path = hf_hub_download(
+            repo_id="nihardon/fine-tuned-unit-test-generator",
+            filename="llama-3-8b.Q4_K_M.gguf",
+        )
+        
+        load_status = "🧠 Loading into RAM (approx 60s)..."
+        print("Loading weights...")
+        
+        # Load the model (verbose=False speeds it up slightly)
+        llm_model = Llama(
+            model_path=model_path,
+            n_ctx=1024,
+            n_threads=2,
+            verbose=False 
+        )
+        
+        load_status = "✅ Model Ready!"
+        is_loaded = True
+        print("🚀 Model successfully loaded!")
+        
+    except Exception as e:
+        load_status = f"❌ Error: {str(e)}"
+        print(load_status)
 
-    print("⏳ First-time load: Downloading/Loading model...")
-    
-    # 1. Download (Cached)
-    model_path = hf_hub_download(
-        repo_id="nihardon/fine-tuned-unit-test-generator",
-        filename="llama-3-8b.Q4_K_M.gguf",
-    )
-    
-    # 2. Load into RAM
-    llm_model = Llama(
-        model_path=model_path,
-        n_ctx=1024,        # Context window
-        n_threads=2,       # Use 2 threads for better speed
-        verbose=False,     # Reduce logs to prevent buffer lag
-    )
-    
-    print("✅ Model loaded!")
-    return llm_model
+# Start the loader immediately in the background
+threading.Thread(target=load_model_in_background, daemon=True).start()
 
+# --- The Generator Function ---
 def generate_test(user_code):
-    # Load model (only happens once)
-    llm = load_model()
+    global llm_model, is_loaded, load_status
+    
+    # 1. Check if model is ready
+    if not is_loaded or llm_model is None:
+        return f"⚠️ SYSTEM INITIALIZING...\n\nCurrent Status: {load_status}\n\nPlease wait 60 seconds and click Generate again."
     
+    # 2. Run Generation
     prompt = f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
 
 ### Instruction:
@@ -46,19 +64,21 @@ You are an expert Python QA engineer. Write a pytest unit test for the following
 
 ### Response:
 """
-    output = llm(
-        prompt,
-        max_tokens=512,
-        stop=["### Instruction:", "### Input:"],
-        echo=False
-    )
-    return output["choices"][0]["text"].strip()
+    try:
+        output = llm_model(
+            prompt,
+            max_tokens=512,
+            stop=["### Instruction:", "### Input:"],
+            echo=False
+        )
+        return output["choices"][0]["text"].strip()
+    except Exception as e:
+        return f"Error during generation: {str(e)}"
 
 # --- The UI ---
-# This part runs instantly, so the Health Check passes immediately!
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🧪 AI Unit Test Generator")
-    gr.Markdown("*Note: The first request will take ~1 minute to initialize the model.*")
+    gr.Markdown("**Status:** System starts automatically. If the model isn't ready, it will tell you.")
     
     with gr.Row():
         with gr.Column():
@@ -73,6 +93,4 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             
     btn.click(generate_test, inputs=input_box, outputs=output_box)
 
-# Launch
-print("🚀 Server starting...")
 demo.launch(server_name="0.0.0.0", server_port=7860)