Update app.py

app.py

@@ -1,9 +1,9 @@
 import os
 import glob
 import json
-import psutil # Added to check system health
+import psutil
 from fastapi import FastAPI, Request, HTTPException
-from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi.responses import StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
 from llama_cpp import Llama
 
@@ -17,17 +17,21 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# --- Config ---
+# --- Configuration ---
+# Map the real filenames to your preferred names
+MODEL_MAP = {
+    "qwen2.5-0.5b-instruct-q2_k.gguf": "Hannah-1.0 Light",
+    "qwen2.5-0.5b-instruct-q4_k_m.gguf": "Hannah-1.0 Heavy"
+}
+
 current_model = None
 current_model_name = ""
 
 def get_model(model_name):
     global current_model, current_model_name
-    if not model_name: raise HTTPException(status_code=400, detail="No model selected")
     
-    # Check if file actually exists
-    if not os.path.exists(model_name):
-        raise HTTPException(status_code=404, detail=f"Model file {model_name} not found inside Space.")
+    if not model_name: raise HTTPException(status_code=400, detail="No model selected")
+    if not os.path.exists(model_name): raise HTTPException(status_code=404, detail="Model file not found")
 
     if current_model_name == model_name and current_model is not None:
         return current_model
@@ -35,12 +39,12 @@ def get_model(model_name):
     print(f"Loading {model_name}...")
     if current_model is not None: del current_model
     
-    # Optimized for < 1GB models
+    # Speed Optimization for 0.5B
     current_model = Llama(
         model_path=model_name,
-        n_ctx=4096,     # High context window
-        n_threads=2,    # Free Tier Max
-        n_batch=1024,
+        n_ctx=4096,     
+        n_threads=2,    
+        n_batch=1024,   
         verbose=False
     )
     current_model_name = model_name
@@ -48,20 +52,23 @@ def get_model(model_name):
 
 @app.get("/api/models")
 async def list_models():
-    # Returns file size and name for the table
-    models = []
+    models_info = []
+    # Only look for the files you uploaded
     for f in glob.glob("*.gguf"):
+        display_name = MODEL_MAP.get(f, f) # Use custom name if available, else filename
         size_mb = os.path.getsize(f) / (1024 * 1024)
-        models.append({"name": f, "size": f"{size_mb:.1f} MB"})
-    return {"models": models}
+        models_info.append({
+            "filename": f, 
+            "display_name": display_name,
+            "size": f"{size_mb:.1f} MB"
+        })
+    return {"models": models_info}
 
 @app.get("/api/status")
 async def system_status():
-    # Helper to show RAM usage in the table
     ram = psutil.virtual_memory()
     return {
         "ram_used": f"{ram.used / (1024*1024):.0f} MB",
-        "ram_total": f"{ram.total / (1024*1024):.0f} MB",
         "cpu": f"{psutil.cpu_percent()}%"
     }
 
@@ -84,6 +91,7 @@ async def chat(request: Request):
     llm = get_model(model_file)
 
     def iter_response():
+        # Standard ChatML Prompt
         prompt = f"<|im_start|>system\nYou are Hannah 1.0, an intelligent pilot assistant.<|im_end|>\n<|im_start|>user\n{user_input}<|im_end|>\n<|im_start|>assistant\n"
         stream = llm(prompt, max_tokens=2048, stop=["<|im_end|>"], stream=True)
         for output in stream: