v3

app.py

@@ -1,9 +1,9 @@
-from starlette.types import Message
 from fastapi import FastAPI, HTTPException
 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")
@@ -12,8 +12,6 @@ client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 app = FastAPI()
 
 # Modèle pour les données d'entrée
-
-
 class PredictionRequest(BaseModel):
     message: str
     history: List[Tuple[str, str]] = []
@@ -23,9 +21,10 @@ class PredictionRequest(BaseModel):
     top_p: float = 0.95
 
 
-def generate_stream(request: PredictionRequest):
+async def generate_stream(request: PredictionRequest):
     """
-    Générateur synchrone pour produire les tokens progressivement.
+    Générateur asynchrone pour produire les tokens progressivement.
+    Utilise asyncio.to_thread pour rendre l'appel synchrone compatible avec async.
     """
     messages = [{"role": "system", "content": request.system_message}]
     for user_input, assistant_response in request.history:
@@ -36,57 +35,35 @@ def generate_stream(request: PredictionRequest):
                 {"role": "assistant", "content": assistant_response})
     messages.append({"role": "user", "content": request.message})
 
-    yield "START\n".encode("utf-8")
-
     try:
-        # Appel à l'API Hugging Face avec streaming
-        for message in client.chat_completion(
-            messages,
-            max_tokens=request.max_tokens,
-            stream=True,
-            temperature=request.temperature,
-            top_p=request.top_p,
-        ):
+        # Exécution du client synchrone dans un thread séparé
+        def sync_stream():
+            return client.chat_completion(
+                messages,
+                max_tokens=request.max_tokens,
+                stream=True,
+                temperature=request.temperature,
+                top_p=request.top_p,
+            )
+
+        # Appel synchrone dans un thread asynchrone
+        for message in await asyncio.to_thread(sync_stream):
             token = message.choices[0].delta.content
-            print(token)
-            # Chaque token avec un saut de ligne
-            yield f"{token}\n".encode("utf-8")
+            yield f"{token}\n"
     except Exception as e:
-        yield f"Error: {str(e)}\n".encode("utf-8")
-
-
-class CustomStreamingResponse(StreamingResponse):
-    """
-    Personnalisation de StreamingResponse pour s'assurer que chaque chunk est envoyé immédiatement.
-    """
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.started = False  # Initialisation de l'attribut `started`
-
-    async def stream_response(self, send: Message):
-        # Envoi du message de démarrage une seule fois
-        if not self.started:
-            await send({"type": "http.response.start", "status": 200, "headers": [(b"content-type", b"text/plain")]})
-            self.started = True
-
-        # Envoi des chunks de réponse
-        async for chunk in self.body_iterator:
-            await send({"type": "http.response.body", "body": chunk, "more_body": True})
-        await send({"type": "http.response.body", "body": b"", "more_body": False})
+        yield f"Error: {str(e)}\n"
 
 
 @app.post("/predict")
-def predict(request: PredictionRequest):
+async def predict(request: PredictionRequest):
     """
-    Endpoint REST avec réponse en streaming synchrone.
+    Endpoint REST avec réponse en streaming.
     """
-    return CustomStreamingResponse(
+    return StreamingResponse(
         generate_stream(request),
-        media_type="text/plain"  # Peut être changé en JSON si nécessaire
+        media_type="text/plain"  # Peut être changé en JSON si besoin
     )
 
-
 # Pour le test en local
 if __name__ == "__main__":
     import uvicorn

client.py

@@ -5,10 +5,10 @@ url = "https://micksoftware-laria-startup.hf.space/predict"
 
 # Données pour la requête
 payload = {
-    "message": "racontes moi une histoire de 50 mots",
+    "message": "quelle longueur d'ypothenuse pour un triangle de cote 4 et 9",
     "history": [],
     "system_message": "You are a friendly Chatbot.",
-    "max_tokens": 2048,
+    "max_tokens": 512,
     "temperature": 0.7,
     "top_p": 0.95,
 }
@@ -21,6 +21,6 @@ if response.status_code == 200:
     print("Streaming response:")
     for chunk in response.iter_lines(decode_unicode=True):
         if chunk:
-            print(chunk, end="\n")
+            print(chunk, end="")
 else:
     print(f"Erreur : {response.status_code} - {response.text}")