Implement request batching

app.py

@@ -10,10 +10,12 @@ from typing import TypedDict, Self, Any
 import json
 from pathlib import Path
 from dataclasses import dataclass
+from itertools import batched, chain
 
 MODEL_NAME = "all-MiniLM-L6-v2"  # TODO: make configurable
 DIR = Path("index")
 SEARCH_TIME_S = 1  # TODO: optimize
+K = 20
 
 
 class IndexParameters(TypedDict):
@@ -119,12 +121,15 @@ def get_index(dir: Path, search_time_s: float) -> Dataset:
 
 
 def execute_request(ids: list[str]) -> list[Work]:
+    if len(ids) > 100:
+        raise ValueError("querying /works endpoint with more than 100 works")
+
     # query with the /works endpoint with a specific list of IDs and fields
     search_filter = f"openalex_id:{"|".join(ids)}"
     search_select = ",".join(["id"] + Work.get_raw_fields())
     response = requests.get(
         "https://api.openalex.org/works",
-        {"filter": search_filter, "select": search_select}
+        {"filter": search_filter, "select": search_select, "per-page": 100}
     )
     response.raise_for_status()
 
@@ -196,17 +201,19 @@ model = get_model(MODEL_NAME, "cpu")
 index = get_index(DIR, SEARCH_TIME_S)
 
 
-def search(query: str) -> str:
-    global model, index
-
+# function signature: (expanded tuple of input batches) -> tuple of output batches
+def search(query: list[str]) -> tuple[list[str]]:
     query_embedding = model.encode(query)
-    distances, faiss_ids = index.search("embeddings", query_embedding, 20)
-    openalex_ids = index[faiss_ids]["idxs"]
+    distances, faiss_ids = index.search_batch("embeddings", query_embedding, K)
 
-    works = execute_request(openalex_ids)
-    result_string = format_response(works, distances)
+    faiss_ids_flat = list(chain(*faiss_ids))
+    openalex_ids_flat = index[faiss_ids_flat]["idxs"]
+    works_flat = execute_request(openalex_ids_flat)
+    works = [list(batch) for batch in batched(works_flat, K)]
 
-    return result_string
+    result_strings = [format_response(w, d) for w, d in zip(works, distances)]
+
+    return (result_strings, )
 
 
 with gr.Blocks() as demo:
@@ -233,8 +240,8 @@ with gr.Blocks() as demo:
         container=True,
     )
 
-    query.submit(search, inputs=[query], outputs=[results])
-    btn.click(search, inputs=[query], outputs=[results])
+    query.submit(search, inputs=[query], outputs=[results], batch=True)
+    btn.click(search, inputs=[query], outputs=[results], batch=True)
 
-demo.queue(2)
+demo.queue()
 demo.launch()