streaming

app.py

@@ -1,7 +1,9 @@
 from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
+from fastapi.responses import StreamingResponse
 from huggingface_hub import InferenceClient
+from pydantic import BaseModel
 from typing import List, Tuple
+import asyncio
 
 # Initialisation du client Hugging Face
 client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
@@ -18,12 +20,11 @@ class PredictionRequest(BaseModel):
     temperature: float = 0.7
     top_p: float = 0.95
 
-@app.post("/predict")
-async def predict(request: PredictionRequest):
+
+async def generate_stream(request: PredictionRequest):
     """
-    Endpoint REST pour effectuer une prédiction.
+    Générateur asynchrone pour le streaming de réponse.
     """
-    # Préparer les messages pour l'inférence
     messages = [{"role": "system", "content": request.system_message}]
     for user_input, assistant_response in request.history:
         if user_input:
@@ -32,21 +33,30 @@ async def predict(request: PredictionRequest):
             messages.append({"role": "assistant", "content": assistant_response})
     messages.append({"role": "user", "content": request.message})
     
-    # Appel de l'API Hugging Face
     try:
-        response = ""
-        for message in client.chat_completion(
+        async for message in client.chat_completion(
             messages,
             max_tokens=request.max_tokens,
             stream=True,
             temperature=request.temperature,
             top_p=request.top_p,
         ):
-            response += message.choices[0].delta.content
-        return {"response": response}
+            token = message.choices[0].delta.content
+            yield token
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+
+@app.post("/predict")
+async def predict(request: PredictionRequest):
+    """
+    Endpoint REST avec réponse en streaming.
+    """
+    return StreamingResponse(
+        generate_stream(request),
+        media_type="text/plain"  # Peut être changé en JSON si besoin
+    )
+
 # Pour le test en local
 if __name__ == "__main__":
     import uvicorn