precommit

Author: floschne

src/fundus_murag/api/routers/lookup.py

@@ -42,12 +42,8 @@ def list_all_collections():
     summary="Get a `FundusCollection` by its name or MURAG ID.",
 )
 def get_fundus_collection_by_id(
-    collection_name: str | None = Query(
-        None, description="Unique internal name for the collection."
-    ),
-    murag_id: str | None = Query(
-        None, description="Unique identifier for the collection in the VectorDB."
-    ),
+    collection_name: str | None = Query(None, description="Unique internal name for the collection."),
+    murag_id: str | None = Query(None, description="Unique identifier for the collection in the VectorDB."),
 ):
     try:
         if collection_name:
@@ -125,17 +121,13 @@ def get_fundus_records_by_id(
         None,
         description="An identifier for the `FundusRecord`. If a `FundusRecord` has multiple images, the records share the `fundus_id`.",
     ),
-    murag_id: str | None = Query(
-        None, description="Unique identifier for the record in the VectorDB."
-    ),
+    murag_id: str | None = Query(None, description="Unique identifier for the record in the VectorDB."),
 ):
     try:
         if fundus_id is None and murag_id is None:
             raise ValueError("Either `fundus_id` or `murag_id` must be provided.")
         elif murag_id is not None and fundus_id is not None:
-            raise ValueError(
-                "Either `fundus_id` or `murag_id` must be provided, not both."
-            )
+            raise ValueError("Either `fundus_id` or `murag_id` must be provided, not both.")
         elif murag_id:
             record = vdb.get_fundus_record_by_murag_id(murag_id=murag_id)
         elif fundus_id:
@@ -159,9 +151,7 @@ def get_fundus_records_by_id(
     summary="Returns the `FundusRecordImage`s from the `FundusRecord` with the specified `murag_id`.",
 )
 def get_fundus_record_image_by_murag_id(
-    murag_id: str = Query(
-        ..., description="Unique identifier for the `FundusRecord` in the VectorDB."
-    ),
+    murag_id: str = Query(..., description="Unique identifier for the `FundusRecord` in the VectorDB."),
 ):
     try:
         record_img = vdb.get_fundus_record_image_by_murag_id(murag_id=murag_id)

src/fundus_murag/api/routers/random.py

@@ -35,9 +35,7 @@ def get_random_fundus_collection(
 )
 def get_random_fundus_record(
     n: int = Query(1, description="The number of random records to return."),
-    collection_name: str | None = Query(
-        None, description="Unique internal name for the collection."
-    ),
+    collection_name: str | None = Query(None, description="Unique internal name for the collection."),
 ):
     try:
         return vdb.get_random_fundus_records(n=n, collection_name=collection_name)

src/fundus_murag/api/routers/search.py

@@ -29,9 +29,7 @@
 )
 def fundus_record_i2i_similarity_search(query: SimilaritySearchQuery):
     try:
-        query_embedding = mlc.compute_image_embedding(
-            base64_image=query.query, return_tensor="np"
-        ).tolist()  # type: ignore
+        query_embedding = mlc.compute_image_embedding(base64_image=query.query, return_tensor="np").tolist()  # type: ignore
         return vdb._fundus_record_image_similarity_search(
             query_embedding=query_embedding,
             search_in_collections=query.collection_names,
@@ -48,9 +46,7 @@ def fundus_record_i2i_similarity_search(query: SimilaritySearchQuery):
 )
 def fundus_record_t2i_similarity_search(query: SimilaritySearchQuery):
     try:
-        query_embedding = mlc.compute_text_embedding(
-            text=query.query, return_tensor="np"
-        ).tolist()  # type: ignore
+        query_embedding = mlc.compute_text_embedding(text=query.query, return_tensor="np").tolist()  # type: ignore
         return vdb._fundus_record_image_similarity_search(
             query_embedding=query_embedding,
             search_in_collections=query.collection_names,
@@ -67,9 +63,7 @@ def fundus_record_t2i_similarity_search(query: SimilaritySearchQuery):
 )
 def fundus_record_i2t_similarity_search(query: SimilaritySearchQuery):
     try:
-        query_embedding = mlc.compute_image_embedding(
-            base64_image=query.query, return_tensor="np"
-        ).tolist()  # type: ignore
+        query_embedding = mlc.compute_image_embedding(base64_image=query.query, return_tensor="np").tolist()  # type: ignore
         return vdb._fundus_record_title_similarity_search(
             query_embedding=query_embedding,
             search_in_collections=query.collection_names,
@@ -86,9 +80,7 @@ def fundus_record_i2t_similarity_search(query: SimilaritySearchQuery):
 )
 def fundus_record_t2t_similarity_search(query: SimilaritySearchQuery):
     try:
-        query_embedding = mlc.compute_text_embedding(
-            text=query.query, return_tensor="np"
-        ).tolist()  # type: ignore
+        query_embedding = mlc.compute_text_embedding(text=query.query, return_tensor="np").tolist()  # type: ignore
         return vdb._fundus_record_title_similarity_search(
             query_embedding=query_embedding,
             search_in_collections=query.collection_names,
@@ -140,9 +132,7 @@ def fundus_collection_lexical_search(query: CollectionLexicalSearchQuery):
     summary="Perform a semantic similarity search on `FundusCollection`s based on their description.",
 )
 def fundus_collection_description_similarity_search(query: SimilaritySearchQuery):
-    query_embedding = mlc.compute_text_embedding(
-        text=query.query, return_tensor="np"
-    ).tolist()  # type: ignore
+    query_embedding = mlc.compute_text_embedding(text=query.query, return_tensor="np").tolist()  # type: ignore
     try:
         return vdb.fundus_collection_description_similarity_search(
             query_embedding=query_embedding,
@@ -159,9 +149,7 @@ def fundus_collection_description_similarity_search(query: SimilaritySearchQuery
 )
 def fundus_collection_title_similarity_search(query: SimilaritySearchQuery):
     try:
-        query_embedding = mlc.compute_text_embedding(
-            text=query.query, return_tensor="np"
-        ).tolist()  # type: ignore
+        query_embedding = mlc.compute_text_embedding(text=query.query, return_tensor="np").tolist()  # type: ignore
         return vdb.fundus_collection_title_similarity_search(
             query_embedding=query_embedding,
             top_k=query.top_k,

src/fundus_murag/ml/client.py

@@ -29,9 +29,7 @@ def _wait_for_ready(self, s: int = 60, sleep_t: int = 3) -> None:
             if self._is_ready():
                 logger.info(f"Fundus ML is ready at {self._fundus_ml_url}!")
                 return
-            logger.info(
-                f"Waiting {sleep_t}s for Fundus ML to be ready at {self._fundus_ml_url}..."
-            )
+            logger.info(f"Waiting {sleep_t}s for Fundus ML to be ready at {self._fundus_ml_url}...")
             time.sleep(sleep_t)
             s -= sleep_t
 
@@ -85,9 +83,7 @@ def _get_embeddings(
         return_tensor: Literal["pt", "np"] | None = "np",
         squeeze: bool = True,
     ) -> "EmbeddingsOutput | np.ndarray | torch.Tensor":
-        response = requests.post(
-            f"{self._fundus_ml_url}/embed", json=input.model_dump()
-        )
+        response = requests.post(f"{self._fundus_ml_url}/embed", json=input.model_dump())
         response.raise_for_status()
         response_json = response.json()
         emb = EmbeddingsOutput.model_validate(response_json)
@@ -103,11 +99,7 @@ def _get_embeddings(
             if squeeze:
                 emb = emb.squeeze()
         else:
-            if (
-                squeeze
-                and len(emb.embeddings) == 1
-                and isinstance(emb.embeddings[0], list)
-            ):
+            if squeeze and len(emb.embeddings) == 1 and isinstance(emb.embeddings[0], list):
                 emb.embeddings = emb.embeddings[0]
 
         return emb

src/fundus_murag/ml/server.py

@@ -37,21 +37,11 @@ def setup(self, device: str):
 
     def decode_request(self, request: EmbeddingsInput) -> dict[str, list | None]:
         if request.input_type == "text":
-            text = (
-                request.input_data
-                if isinstance(request.input_data, list)
-                else [request.input_data]
-            )
+            text = request.input_data if isinstance(request.input_data, list) else [request.input_data]
             image = None
         elif request.input_type == "image":
-            image_data = (
-                request.input_data
-                if isinstance(request.input_data, list)
-                else [request.input_data]
-            )
-            image = [
-                Image.open(io.BytesIO(base64.b64decode(b64))) for b64 in image_data
-            ]
+            image_data = request.input_data if isinstance(request.input_data, list) else [request.input_data]
+            image = [Image.open(io.BytesIO(base64.b64decode(b64))) for b64 in image_data]
             text = None
         else:
             raise ValueError("Invalid request type")
@@ -65,9 +55,7 @@ def _compute_text_embedding(self, text_features: BatchFeature) -> torch.Tensor:
 
     def _compute_image_embedding(self, image_features: BatchFeature) -> torch.Tensor:
         with torch.no_grad():
-            img_emb = self.model.get_image_features(
-                **image_features.to(self.device, dtype=TORCH_DTYPE)
-            )
+            img_emb = self.model.get_image_features(**image_features.to(self.device, dtype=TORCH_DTYPE))
         return img_emb
 
     def predict(self, inputs: dict[str, list | None]) -> torch.Tensor: