optimized code for tread safe operations

main.py

@@ -1,3 +1,4 @@
+import os
 import uuid
 import logging
 import asyncio
@@ -25,6 +26,7 @@ MAX_BATCH_SIZE = 128
 # ----------------------------- #
 #        Shared Storage         #
 # ----------------------------- #
+lock = asyncio.Lock()
 query_queue: asyncio.Queue = asyncio.Queue()
 results: dict[str, dict] = {}
 classifier = None  # will be initialized in lifespan
@@ -56,6 +58,8 @@ def load_classifier(model_name: str):
         model=model,
         tokenizer=tokenizer,
         framework="pt",
+        batch_size=MAX_BATCH_SIZE,
+        num_workers=os.cpu_count(),
     )
 
 
@@ -82,14 +86,19 @@ async def process_queue():
 
         sentences = [item["sentence"] for item in batch]
         ids = [item["id"] for item in batch]
-        predictions = classifier(sentences, batch_size=len(sentences))
+        predictions = classifier(sentences)
 
-        for query_id, pred, sentence in zip(ids, predictions, sentences):
-            results[query_id] = {
-                "sentence": sentence,
-                "label": pred["label"],
-                "score": pred["score"],
-            }
+        async with lock:
+            results.update(
+                {
+                    query_id: {
+                        "sentence": sentence,
+                        "label": pred["label"],
+                        "score": pred["score"],
+                    }
+                    for query_id, pred, sentence in zip(ids, predictions, sentences)
+                }
+            )
 
 
 # ----------------------------- #
@@ -127,10 +136,14 @@ async def classify(query: Query):
     query_id = str(uuid.uuid4())
     await query_queue.put({"id": query_id, "sentence": query.sentence})
 
-    while query_id not in results:
+    result = None
+    while result is None:
+        async with lock:
+            if query_id in results:
+                result = results.pop(query_id)
         await asyncio.sleep(0.001)
 
-    return {"id": query_id, "result": results.pop(query_id)}
+    return {"id": query_id, "result": result}
 
 
 app.mount("/", StaticFiles(directory="static", html=True), name="static")