Merge branch 'main' of https://github.com/yichuan-w/LEANN

packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_embedding_server.py

@@ -57,7 +57,6 @@ def create_hnsw_embedding_server(
     finally:
         sys.path.pop(0)
 
-
     # Check port availability
     import socket
 
@@ -152,7 +151,9 @@ def create_hnsw_embedding_server(
                             raise
 
                     # Process embeddings
-                    embeddings = compute_embeddings(texts, model_name, mode=embedding_mode)
+                    embeddings = compute_embeddings(
+                        texts, model_name, mode=embedding_mode
+                    )
                     print(
                         f"INFO: Computed embeddings for {len(texts)} texts, shape: {embeddings.shape}"
                     )
@@ -278,7 +279,7 @@ if __name__ == "__main__":
         "--embedding-mode",
         type=str,
         default="sentence-transformers",
-        choices=["sentence-transformers", "openai"],
+        choices=["sentence-transformers", "openai", "mlx"],
         help="Embedding backend mode",
     )
 

packages/leann-core/src/leann/searcher_base.py

@@ -42,8 +42,10 @@ class BaseSearcher(LeannBackendSearcherInterface, ABC):
                 "WARNING: embedding_model not found in meta.json. Recompute will fail."
             )
 
+        self.embedding_mode = self.meta.get("embedding_mode", "sentence-transformers")
+
         self.embedding_server_manager = EmbeddingServerManager(
-            backend_module_name=backend_module_name
+            backend_module_name=backend_module_name,
         )
 
     def _load_meta(self) -> Dict[str, Any]:
@@ -67,14 +69,12 @@ class BaseSearcher(LeannBackendSearcherInterface, ABC):
                 "Cannot use recompute mode without 'embedding_model' in meta.json."
             )
 
-        embedding_mode = self.meta.get("embedding_mode", "sentence-transformers")
-
         server_started, actual_port = self.embedding_server_manager.start_server(
             port=port,
             model_name=self.embedding_model,
+            embedding_mode=self.embedding_mode,
             passages_file=passages_source_file,
             distance_metric=kwargs.get("distance_metric"),
-            embedding_mode=embedding_mode,
             enable_warmup=kwargs.get("enable_warmup", False),
         )
         if not server_started:

test/build_mlx_index.py

@@ -12,7 +12,7 @@ else:
     builder = LeannBuilder(
         backend_name="hnsw",
         embedding_model="mlx-community/Qwen3-Embedding-0.6B-4bit-DWQ",
-        use_mlx=True,
+        embedding_mode="mlx",
     )
 
     # 2. Add documents