xx

app.py

@@ -1,22 +1,29 @@
-from fastapi import FastAPI
+import os
+from fastapi import FastAPI, Request, HTTPException, Header
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
 app = FastAPI()
 
+# Récupérer le token depuis les variables d’environnement
+API_TOKEN = os.environ.get("API_TOKEN")
+
 # Charger le modèle
 model_name = "google/medgemma-4b-pt"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
 
-# Requête attendue
-class Request(BaseModel):
+# Modèle de requête
+class GenerationRequest(BaseModel):
     prompt: str
 
 @app.post("/generate")
-def generate(request: Request):
-    inputs = tokenizer(request.prompt, return_tensors="pt").to(model.device)
+async def generate(request_data: GenerationRequest, authorization: str = Header(None)):
+    if authorization != f"Bearer {API_TOKEN}":
+        raise HTTPException(status_code=401, detail="Unauthorized")
+
+    inputs = tokenizer(request_data.prompt, return_tensors="pt").to(model.device)
     with torch.no_grad():
         outputs = model.generate(**inputs, max_new_tokens=100)
     result = tokenizer.decode(outputs[0], skip_special_tokens=True)