Update main.py

main.py

@@ -1,8 +1,9 @@
 from fastapi import FastAPI
+from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 from huggingface_hub import InferenceClient
 import uvicorn
-
+import asyncio
 
 app = FastAPI()
 
@@ -25,10 +26,8 @@ def format_prompt(message, history):
     prompt += f"[INST] {message} [/INST]"
     return prompt
 
-def generate(item: Item):
-    temperature = float(item.temperature)
-    if temperature < 1e-2:
-        temperature = 1e-2
+async def generate(item: Item):
+    temperature = max(float(item.temperature), 1e-2)  # Ensure temperature is not too low
     top_p = float(item.top_p)
 
     generate_kwargs = dict(
@@ -41,14 +40,19 @@ def generate(item: Item):
     )
 
     formatted_prompt = format_prompt(f"{item.system_prompt}, {item.prompt}", item.history)
-    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
-    output = ""
-
-    for response in stream:
-        output += response.token.text
-    return output
+    
+    # Stream the response from the model
+    async def event_stream():
+        stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
+        
+        async for response in stream:
+            yield response.token.text  # Yield each token as it is received
+            
+            # Optional: Add a small delay to simulate streaming effect (if needed)
+            await asyncio.sleep(0.1)
+
+    return event_stream()
 
 @app.post("/generate/")
 async def generate_text(item: Item):
-    return {"response": generate(item)}
-
+    return StreamingResponse(generate(item), media_type="text/event-stream")