Update app.py

app.py

@@ -5,6 +5,9 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 
 app = FastAPI()
 
+# -------------------------------------
+# MODEL (FAST & SMALL)
+# -------------------------------------
 MODEL_REPO = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 
 tokenizer = None
@@ -14,25 +17,29 @@ model = None
 def load_model():
     global tokenizer, model
     if tokenizer is None or model is None:
-        print("🔥 Loading model now (lazy load)... This will take time but only once.")
+        print("🔥 Loading TinyLlama model...")
 
         tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
 
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_REPO,
-            dtype=torch.float16,
-            device_map="cpu",   # force CPU for Spaces
+            torch_dtype=torch.float32,    # CPU safe
+            device_map="cpu",
             low_cpu_mem_usage=True
         )
 
-        print("✅ Model loaded successfully!")
+        print("✅ TinyLlama loaded successfully!")
 
 
+# -------------------------------------
+# ROUTES
+# -------------------------------------
 @app.get("/")
-async def root():
+async def home():
     return {
-        "message": "🚀 FastAPI MPT Model Running on Hugging Face Spaces",
-        "endpoints": ["/", "/status", "/generate"]
+        "message": "🚀 TinyLlama Chat API (FastAPI + HF Spaces)",
+        "endpoints": ["/", "/status", "/generate"],
+        "model": MODEL_REPO
     }
 
 
@@ -41,7 +48,7 @@ async def status():
     return {
         "status": "ok",
         "model": MODEL_REPO,
-        "model_loaded": model is not None
+        "loaded": model is not None
     }
 
 
@@ -51,17 +58,19 @@ class InputText(BaseModel):
 
 @app.post("/generate")
 async def generate_text(data: InputText):
-
-    # Load model ONLY when first request happens
     load_model()
 
-    inputs = tokenizer(data.text, return_tensors="pt").to(model.device)
+    prompt = f"<|system|>You are a friendly helpful AI assistant.<|user|>{data.text}<|assistant|>"
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
     output = model.generate(
         **inputs,
         max_new_tokens=150,
-        temperature=0.7
+        temperature=0.7,
+        top_p=0.9,
+        do_sample=True
     )
 
-    text = tokenizer.decode(output[0], skip_special_tokens=True)
-    return {"response": text}
+    result = tokenizer.decode(output[0], skip_special_tokens=True)
+    return {"response": result}