Update app.py

app.py

@@ -5,9 +5,6 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 
 app = FastAPI()
 
-# -------------------------------------
-# MODEL (FAST & SMALL)
-# -------------------------------------
 MODEL_REPO = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 
 tokenizer = None
@@ -16,28 +13,27 @@ model = None
 
 def load_model():
     global tokenizer, model
+
     if tokenizer is None or model is None:
-        print("🔥 Loading TinyLlama model...")
+        print("🔥 Loading TinyLlama model on CPU...")
 
         tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
 
+        # ❗ NO device_map
+        # ❗ NO torch_dtype=float16
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_REPO,
-            torch_dtype=torch.float32,    # CPU safe
-            device_map="cpu",
+            torch_dtype=torch.float32,      # safe CPU
             low_cpu_mem_usage=True
         )
 
-        print("✅ TinyLlama loaded successfully!")
+        print("✅ TinyLlama loaded!")
 
 
-# -------------------------------------
-# ROUTES
-# -------------------------------------
 @app.get("/")
 async def home():
     return {
-        "message": "🚀 TinyLlama Chat API (FastAPI + HF Spaces)",
+        "message": "🚀 TinyLlama Chat API Running",
         "endpoints": ["/", "/status", "/generate"],
         "model": MODEL_REPO
     }
@@ -60,9 +56,13 @@ class InputText(BaseModel):
 async def generate_text(data: InputText):
     load_model()
 
-    prompt = f"<|system|>You are a friendly helpful AI assistant.<|user|>{data.text}<|assistant|>"
+    prompt = f"<|system|>You are a helpful assistant.<|user|>{data.text}<|assistant|>"
+
+    inputs = tokenizer(prompt, return_tensors="pt")
 
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    # Move to CPU explicitly
+    inputs = {k: v.to("cpu") for k, v in inputs.items()}
+    model.to("cpu")
 
     output = model.generate(
         **inputs,