Update app.py

app.py

@@ -1,24 +1,41 @@
-print("starting")
-
 import gradio as gr
 from huggingface_hub import hf_hub_download
 from llama_cpp import Llama
 
-model_path = hf_hub_download(
-    repo_id="nihardon/fine-tuned-unit-test-generator",
-    filename="llama-3-8b.Q4_K_M.gguf", 
-)
-
-print("downloading model")
-# Load the model (CPU optimized)
-llm = Llama(
-    model_path=model_path,
-    n_ctx=2048,
-    n_threads=2,
-    verbose = True
-)
-print("loading model")
+# Global variable to hold the model (starts empty)
+llm_model = None
+
+def load_model():
+    """Loads the model only when needed."""
+    global llm_model
+    
+    # If already loaded, just return it
+    if llm_model is not None:
+        return llm_model
+
+    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
+
 def generate_test(user_code):
+    # Load model (only happens once)
+    llm = load_model()
+    
     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:
@@ -29,21 +46,33 @@ You are an expert Python QA engineer. Write a pytest unit test for the following
 
 ### Response:
 """
-    output = llm(prompt, max_tokens=256, stop=["### Instruction:"], echo=False)
-    return output['choices'][0]['text'].strip()
+    output = llm(
+        prompt,
+        max_tokens=512,
+        stop=["### Instruction:", "### Input:"],
+        echo=False
+    )
+    return output["choices"][0]["text"].strip()
 
-# UI
+# --- 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("**Model:** Custom Fine-Tuned Llama-3 (GGUF) | **Status:** Running Locally")
+    gr.Markdown("*Note: The first request will take ~1 minute to initialize the model.*")
     
     with gr.Row():
         with gr.Column():
-            input_box = gr.Code(language="python", value="def add(a, b):\n    return a + b", label="Function")
+            input_box = gr.Code(
+                language="python",
+                value="def add(a, b):\n    return a + b",
+                label="Function"
+            )
             btn = gr.Button("Generate Pytest", variant="primary")
         with gr.Column():
             output_box = gr.Code(language="python", label="Generated Test Case")
             
     btn.click(generate_test, inputs=input_box, outputs=output_box)
 
-demo.launch()
+# Launch
+print("🚀 Server starting...")
+demo.launch(server_name="0.0.0.0", server_port=7860)