Update main.py

main.py

@@ -1,13 +1,12 @@
 from fastapi import FastAPI
-from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 from huggingface_hub import InferenceClient
+from fastapi.responses import StreamingResponse
 import uvicorn
-import asyncio
 
 app = FastAPI()
 
-client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.2")
+client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
 
 class Item(BaseModel):
     prompt: str
@@ -26,8 +25,10 @@ def format_prompt(message, history):
     prompt += f"[INST] {message} [/INST]"
     return prompt
 
-async def generate(item: Item):
-    temperature = max(float(item.temperature), 1e-2)  # Ensure temperature is not too low
+async def generate_stream(item: Item):
+    temperature = float(item.temperature)
+    if temperature < 1e-2:
+        temperature = 1e-2
     top_p = float(item.top_p)
 
     generate_kwargs = dict(
@@ -40,19 +41,11 @@ async def generate(item: Item):
     )
 
     formatted_prompt = format_prompt(f"{item.system_prompt}, {item.prompt}", item.history)
-    
-    # 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()
+    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
+
+    for response in stream:
+        yield response.token.text  # Stream each token as it's received
 
 @app.post("/generate/")
 async def generate_text(item: Item):
-    return StreamingResponse(generate(item), media_type="text/event-stream")
+    return StreamingResponse(generate_stream(item), media_type="text/plain")