Update app.py

app.py

@@ -3,7 +3,6 @@ from fastapi.responses import HTMLResponse, StreamingResponse
 from fastapi.staticfiles import StaticFiles
 from modules.pmbl import PMBL
 import torch
-from queue import Queue
 import asyncio
 
 print(f"CUDA available: {torch.cuda.is_available()}")
@@ -17,7 +16,8 @@ app.mount("/static", StaticFiles(directory="static"), name="static")
 app.mount("/templates", StaticFiles(directory="templates"), name="templates")
 
 pmbl = PMBL("./PMB-7b.Q6_K.gguf", gpu_layers=50)
-request_queue = Queue()
+request_queue = asyncio.Queue()
+processing_lock = asyncio.Lock()
 
 @app.head("/")
 @app.get("/")
@@ -26,9 +26,10 @@ def index() -> HTMLResponse:
         return HTMLResponse(content=f.read())
 
 async def process_request(user_input: str, mode: str):
-    history = pmbl.get_chat_history(mode, user_input)
-    async for chunk in pmbl.generate_response(user_input, history, mode):
-        yield chunk
+    async with processing_lock:
+        history = pmbl.get_chat_history(mode, user_input)
+        async for chunk in pmbl.generate_response(user_input, history, mode):
+            yield chunk
 
 @app.post("/chat")
 async def chat(request: Request, background_tasks: BackgroundTasks):
@@ -38,11 +39,8 @@ async def chat(request: Request, background_tasks: BackgroundTasks):
         mode = data["mode"]
 
         async def response_generator():
-            future = asyncio.Future()
-            request_queue.put((future, user_input, mode))
-            await future
-
-            async for chunk in future.result():
+            await request_queue.put((user_input, mode))
+            async for chunk in await process_request(user_input, mode):
                 yield chunk
 
         return StreamingResponse(response_generator(), media_type="text/plain")
@@ -52,11 +50,10 @@ async def chat(request: Request, background_tasks: BackgroundTasks):
 
 async def queue_worker():
     while True:
-        if not request_queue.empty():
-            future, user_input, mode = request_queue.get()
-            result = process_request(user_input, mode)
-            future.set_result(result)
-        await asyncio.sleep(0.1)
+        user_input, mode = await request_queue.get()
+        async for _ in process_request(user_input, mode):
+            pass
+        request_queue.task_done()
 
 @app.on_event("startup")
 async def startup_event():