Merge pull request #152 from presenton/fix/embedding-issue-in-arm

feat(fastapi): reverts to chromadb, saves embedding model

servers/fastapi/chroma/models/onnx/config.json

@@ -0,0 +1,25 @@
+{
+  "_name_or_path": "sentence-transformers/all-MiniLM-L6-v2",
+  "architectures": [
+    "BertModel"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "classifier_dropout": null,
+  "gradient_checkpointing": false,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 384,
+  "initializer_range": 0.02,
+  "intermediate_size": 1536,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 512,
+  "model_type": "bert",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 6,
+  "pad_token_id": 0,
+  "position_embedding_type": "absolute",
+  "transformers_version": "4.27.4",
+  "type_vocab_size": 2,
+  "use_cache": true,
+  "vocab_size": 30522
+}

servers/fastapi/chroma/models/onnx/special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]",
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "unk_token": "[UNK]"
+}

servers/fastapi/chroma/models/onnx/tokenizer_config.json

@@ -0,0 +1,15 @@
+{
+  "cls_token": "[CLS]",
+  "do_basic_tokenize": true,
+  "do_lower_case": true,
+  "mask_token": "[MASK]",
+  "model_max_length": 512,
+  "never_split": null,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "special_tokens_map_file": "/Users/hammad/.cache/huggingface/hub/models--sentence-transformers--all-MiniLM-L6-v2/snapshots/7dbbc90392e2f80f3d3c277d6e90027e55de9125/special_tokens_map.json",
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "BertTokenizer",
+  "unk_token": "[UNK]"
+}

servers/fastapi/qdrant/.lock

@@ -1 +0,0 @@
-tmp lock file

servers/fastapi/qdrant/meta.json

@@ -1 +0,0 @@
-{"collections": {"icons": {"vectors": {"fast-bge-small-en-v1.5": {"size": 384, "distance": "Cosine", "hnsw_config": null, "quantization_config": null, "on_disk": null, "datatype": null, "multivector_config": null}}, "shard_number": null, "sharding_method": null, "replication_factor": null, "write_consistency_factor": null, "on_disk_payload": null, "hnsw_config": null, "wal_config": null, "optimizers_config": null, "init_from": null, "quantization_config": null, "sparse_vectors": null, "strict_mode_config": null}}, "aliases": {}}

servers/fastapi/requirements.txt

@@ -12,6 +12,7 @@ cachetools==5.5.2
 certifi==2025.7.14
 cffi==1.17.1
 charset-normalizer==3.4.2
+chromadb==1.0.15
 click==8.2.1
 coloredlogs==15.0.1
 cryptography==45.0.5
@@ -99,7 +100,6 @@ python-dotenv==1.1.1
 python-multipart==0.0.20
 python-pptx==1.0.2
 PyYAML==6.0.2
-qdrant-client==1.15.0
 redis==6.2.0
 referencing==0.36.2
 requests==2.32.4

servers/fastapi/server.py

@@ -1,8 +1,6 @@
 import uvicorn
 import argparse
 
-from services.icon_finder_service import IconFinderService
-
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Run the FastAPI server")
@@ -14,9 +12,6 @@ if __name__ == "__main__":
     )
     args = parser.parse_args()
 
-    # Initialize the Icons Collection
-    IconFinderService()
-
     uvicorn.run(
         "api.main:app",
         host="0.0.0.0",

servers/fastapi/services/icon_finder_service.py

@@ -1,48 +1,51 @@
 import asyncio
 import json
-from qdrant_client import QdrantClient
+import chromadb
+from chromadb.config import Settings
+from chromadb.utils.embedding_functions import ONNXMiniLM_L6_V2
 
 
 class IconFinderService:
     def __init__(self):
         self.collection_name = "icons"
-        self.client = QdrantClient(path="qdrant")
-        print("Initializing icons collection...")
-        self.client.set_model("BAAI/bge-small-en-v1.5")
+        self.client = chromadb.PersistentClient(
+            path="chroma", settings=Settings(anonymized_telemetry=False)
+        )
         self._initialize_icons_collection()
-        print("Icons collection initialized")
 
     def _initialize_icons_collection(self):
+        self.embedding_function = ONNXMiniLM_L6_V2()
+        self.embedding_function.DOWNLOAD_PATH = "chroma/models"
+        self.embedding_function._download_model_if_not_exists()
         try:
-            self.client.get_collection(self.collection_name)
-        except Exception:
-            self._populate_icons_collection()
-
-    def _populate_icons_collection(self):
-        with open("assets/icons.json", "r") as f:
-            icons = json.load(f)
-
-        documents = []
-        metadata = []
-
-        for each in icons["icons"]:
-            if each["name"].split("-")[-1] == "bold":
-                doc_text = f"{each['name']} {each['tags']}"
-                documents.append(doc_text)
-                metadata.append({"name": each["name"]})
-
-        if documents:
-            self.client.add(
-                collection_name=self.collection_name,
-                documents=documents,
-                metadata=metadata,
+            self.collection = self.client.get_collection(
+                self.collection_name, embedding_function=self.embedding_function
             )
+        except Exception:
+            with open("assets/icons.json", "r") as f:
+                icons = json.load(f)
+
+            documents = []
+            ids = []
+
+            for i, each in enumerate(icons["icons"]):
+                if each["name"].split("-")[-1] == "bold":
+                    doc_text = f"{each['name']} {each['tags']}"
+                    documents.append(doc_text)
+                    ids.append(each["name"])
+
+            if documents:
+                self.collection = self.client.create_collection(
+                    name=self.collection_name,
+                    embedding_function=self.embedding_function,
+                    metadata={"hnsw:space": "cosine"},
+                )
+                self.collection.add(documents=documents, ids=ids)
 
     async def search_icons(self, query: str, k: int = 1):
         result = await asyncio.to_thread(
-            self.client.query,
-            collection_name=self.collection_name,
-            query_text=query,
-            limit=k,
+            self.collection.query,
+            query_texts=[query],
+            n_results=k,
         )
-        return [f"/static/icons/bold/{each.metadata['name']}.png" for each in result]
+        return [f"/static/icons/bold/{each}.png" for each in result["ids"][0]]