Update app.py

app.py

@@ -1,30 +1,52 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
 app = FastAPI()
 
-# Ultra-tiny model (SAFE for free CPU)
+# ✅ Ultra-tiny model (SAFE for free CPU)
 MODEL_NAME = "sshleifer/tiny-gpt2"
 
+# Load tokenizer & model once at startup
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=torch.float32
+)
+model.eval()
 
+# Request schema
 class Prompt(BaseModel):
     message: str
 
+# Health check
 @app.get("/")
 def root():
     return {"status": "TinyLLM API is running"}
 
+# Chat endpoint
 @app.post("/chat")
 def chat(prompt: Prompt):
-    inputs = tokenizer(prompt.message, return_tensors="pt")
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=50,
-        do_sample=True,
-        temperature=0.7
+    inputs = tokenizer(
+        prompt.message,
+        return_tensors="pt",
+        truncation=True,
+        max_length=128
     )
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=50,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9
+        )
+
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return {"response": response}
+
+    return {
+        "input": prompt.message,
+        "response": response
+    }