Update app.py

app.py

@@ -1,52 +1,77 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, Header, HTTPException
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+from typing import Optional
 
 app = FastAPI()
 
-# ✅ Ultra-tiny model (SAFE for free CPU)
-MODEL_NAME = "sshleifer/tiny-gpt2"
+# 🔐 CHANGE THIS TO YOUR REAL SECRET KEY
+API_KEY = "sk-tinyllm-9f3a2c7e8b4d1a6c0e52f91d"
+
+# ✅ Best FREE CPU chat model
+MODEL_NAME = "Qwen/Qwen1.5-0.5B-Chat"
 
-# Load tokenizer & model once at startup
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
-    torch_dtype=torch.float32
+    torch_dtype=torch.float32,
+    device_map="cpu"
 )
 model.eval()
 
-# Request schema
 class Prompt(BaseModel):
     message: str
 
-# Health check
+# 🔐 API KEY CHECK (OpenAI style)
+def check_api_key(authorization: Optional[str]):
+    if authorization is None:
+        raise HTTPException(status_code=401, detail="Missing API key")
+
+    if not authorization.startswith("Bearer "):
+        raise HTTPException(status_code=401, detail="Invalid API key format")
+
+    token = authorization.replace("Bearer ", "").strip()
+    if token != API_KEY:
+        raise HTTPException(status_code=401, detail="Invalid API key")
+
 @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",
-        truncation=True,
-        max_length=128
+def chat(
+    prompt: Prompt,
+    authorization: Optional[str] = Header(None)
+):
+    check_api_key(authorization)
+
+    messages = [
+        {"role": "system", "content": "You are a helpful AI assistant."},
+        {"role": "user", "content": prompt.message}
+    ]
+
+    input_ids = tokenizer.apply_chat_template(
+        messages,
+        return_tensors="pt"
     )
 
     with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=50,
-            do_sample=True,
+        output_ids = model.generate(
+            input_ids,
+            max_new_tokens=200,
             temperature=0.7,
-            top_p=0.9
+            top_p=0.9,
+            repetition_penalty=1.1,
+            do_sample=True
         )
 
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    response = tokenizer.decode(
+        output_ids[0][input_ids.shape[-1]:],
+        skip_special_tokens=True
+    )
 
     return {
-        "input": prompt.message,
-        "response": response
+        "model": MODEL_NAME,
+        "response": response.strip()
     }