From c112956d2dab020d49558cbe0a7e93e41eea9ca8 Mon Sep 17 00:00:00 2001
From: Andy Lee <andylizf@outlook.com>
Date: Mon, 21 Jul 2025 21:29:15 -0700
Subject: [PATCH] fix: mlx

---
 .../leann_backend_hnsw/hnsw_embedding_server.py           | 7 ++++---
 packages/leann-core/src/leann/searcher_base.py            | 8 ++++----
 test/build_mlx_index.py                                   | 2 +-
 3 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_embedding_server.py b/packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_embedding_server.py
index 3d3c9fa..b20959b 100644
--- a/packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_embedding_server.py
+++ b/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}"
                     )
@@ -280,7 +281,7 @@ if __name__ == "__main__":
         "--embedding-mode",
         type=str,
         default="sentence-transformers",
-        choices=["sentence-transformers", "openai"],
+        choices=["sentence-transformers", "openai", "mlx"],
         help="Embedding backend mode",
     )
 
diff --git a/packages/leann-core/src/leann/searcher_base.py b/packages/leann-core/src/leann/searcher_base.py
index 73f979a..8592ed4 100644
--- a/packages/leann-core/src/leann/searcher_base.py
+++ b/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:
diff --git a/test/build_mlx_index.py b/test/build_mlx_index.py
index 502cc24..a76bdee 100644
--- a/test/build_mlx_index.py
+++ b/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