Update app.py

app.py

@@ -2,8 +2,10 @@ from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 import uvicorn
+import os
 
 app = FastAPI()
+API_KEY = os.environ.get("API_KEY")
 
 try:
     model = AutoModelForCausalLM.from_pretrained("petertill/cordia-a6")
@@ -18,7 +20,13 @@ try:
         generated_text: str
 
     @app.post("/generate", response_model=GenerateResponse)
-    async def generate(request: GenerateRequest):
+    async def generate(request: GenerateRequest, authorization: str = Header(None)):
+        if not authorization or not authorization.startswith('Bearer '):
+            raise HTTPException(status_code=401, detail="Unauthorized")
+
+        token = authorization.split('Bearer ')[1]
+        if token != API_KEY:
+            raise HTTPException(status_code=401, detail="Unauthorized")
         try:
             output = pipe(request.prompt, max_length=200)[0]['generated_text']
             return GenerateResponse(generated_text=output)