CoolFace
Apppublic

malepati/custom_template_working

sourceHugging Faceupdated 8mo agoView on Hugging Face
0likes
icon_finder_service.py62 linesDownload Raw Back to services
1import asyncio2import json3import chromadb4from chromadb.config import Settings5from chromadb.utils.embedding_functions import ONNXMiniLM_L6_V26 7 8class IconFinderService:9    def __init__(self):10        self.collection_name = "icons"11        self.client = chromadb.PersistentClient(12            path="chroma", settings=Settings(anonymized_telemetry=False)13        )14        print("Initializing icons collection...")15        self._initialize_icons_collection()16        print("Icons collection initialized.")17 18    def _initialize_icons_collection(self):19        self.embedding_function = ONNXMiniLM_L6_V2()20        self.embedding_function.DOWNLOAD_PATH = "chroma/models"21        self.embedding_function._download_model_if_not_exists()22        try:23            self.collection = self.client.get_collection(24                self.collection_name, embedding_function=self.embedding_function25            )26        except Exception:27            with open("assets/icons.json", "r") as f:28                icons = json.load(f)29 30            documents = []31            ids = []32 33            for i, each in enumerate(icons["icons"]):34                if each["name"].split("-")[-1] == "bold":35                    doc_text = f"{each['name']} {each['tags']}"36                    documents.append(doc_text)37                    ids.append(each["name"])38 39            if documents:40                self.collection = self.client.create_collection(41                    name=self.collection_name,42                    embedding_function=self.embedding_function,43                    metadata={"hnsw:space": "cosine"},44                )45                self.collection.add(documents=documents, ids=ids)46 47    async def search_icons(self, query: str, k: int = 1):48        result = await asyncio.to_thread(49            self.collection.query,50            query_texts=[query],51            n_results=k,52        )53        return [f"/static/icons/bold/{each}.svg" for each in result["ids"][0]]54 55 56try:57    ICON_FINDER_SERVICE = IconFinderService()58except Exception as e:59    print(f"Warning: IconFinderService initialization failed: {e}")60    print("Icon search functionality will be disabled.")61    ICON_FINDER_SERVICE = None62