Make random query reproducible

hv0905 · Jun 21, 2024 · 79b5fac · 79b5fac
1 parent 9dba0fc
commit 79b5fac
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 6 deletions.
diff --git a/app/Controllers/search.py b/app/Controllers/search.py
@@ -154,10 +154,14 @@ async def combinedSearch(
 @search_router.get("/random", description="Get random images")
 async def randomPick(
         filter_param: Annotated[FilterParams, Depends(FilterParams)],
-        paging: Annotated[SearchPagingParams, Depends(SearchPagingParams)]) -> SearchApiResponse:
+        paging: Annotated[SearchPagingParams, Depends(SearchPagingParams)],
+        seed: Annotated[int | None, Query(
+            description="The seed for random pick. This is helpful for generating a reproducible random pick.")] = None,
+) -> SearchApiResponse:
     logger.info("Random pick request received")
-    random_vector = services.transformers_service.get_random_vector()
-    result = await services.db_context.querySearch(random_vector, top_k=paging.count, filter_param=filter_param)
+    random_vector = services.transformers_service.get_random_vector(seed)
+    result = await services.db_context.querySearch(random_vector, top_k=paging.count, skip=paging.skip,
+                                                   filter_param=filter_param)
     return await result_postprocessing(
         SearchApiResponse(result=result, message=f"Successfully get {len(result)} results.", query_id=uuid4()))
 

diff --git a/app/Services/transformers_service.py b/app/Services/transformers_service.py
@@ -63,7 +63,7 @@ def get_bert_vector(self, text: str) -> ndarray:
         return vector.cpu().numpy()
 
     @staticmethod
-    def get_random_vector() -> ndarray:
-        vec = np.random.rand(768)
-        vec -= vec.mean()
+    def get_random_vector(seed: int | None = None) -> ndarray:
+        generator = np.random.default_rng(seed)
+        vec = generator.uniform(-1, 1, 768)
         return vec