Update main.py

main.py

@@ -30,7 +30,8 @@ def format_prompt(message, history):
     prompt += f"[INST] {message} [/INST]"
     return prompt
 
-async def generate(prompt, history, temperature=0.9, max_new_tokens=512, top_p=0.95, repetition_penalty=1.0):
+def generate(prompt, history, temperature=0.9, max_new_tokens=512, top_p=0.95, repetition_penalty=1.0):
+    
     temperature = float(temperature)
     if temperature < 1e-2:
         temperature = 1e-2
@@ -47,19 +48,21 @@ async def generate(prompt, history, temperature=0.9, max_new_tokens=512, top_p=0
 
     formatted_prompt = format_prompt(prompt, history)
 
-    async for response in client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False):
-        yield response.token.text
+    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
+        yield output
+    return output
 
 
 
 @app.post("/generate/")
 async def generate_chat(request: Request, prompt: str = Form(...), history: str = Form(...), temperature: float = Form(0.9), max_new_tokens: int = Form(512), top_p: float = Form(0.95), repetition_penalty: float = Form(1.0)):
     history = eval(history)  # Convert history string back to list
-    async def generate_response():
-        async for response_chunk in generate(prompt, history, temperature, max_new_tokens, top_p, repetition_penalty):
-            yield {"response_chunk": response_chunk}
-    return StreamingResponse(generate_response())
+    response = generate(prompt, history, temperature, max_new_tokens, top_p, repetition_penalty)
+    return {"response": response}
 
 app.mount("/", StaticFiles(directory="static", html=True), name="static")