fastembed 0.0.1__tar.gz

fastembed-0.0.1/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

fastembed-0.0.1/PKG-INFO

@@ -0,0 +1,69 @@
+Metadata-Version: 2.1
+Name: fastembed
+Version: 0.0.1
+Summary: Fast Embedding Creation and Simpler API for Qdrant
+Home-page: https://github.com/qdrant/fastembed
+License: Apache License
+Keywords: vector,embedding,neural,search,qdrant,sentence-transformers
+Author: NirantK
+Author-email: nirant.bits@gmail.com
+Requires-Python: >=3.8.0,<3.12
+Classifier: License :: Other/Proprietary License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Requires-Dist: onnxruntime (>=1.15.1,<2.0.0)
+Requires-Dist: onnxruntime-silicon (>=1.15.0,<2.0.0) ; sys_platform == "darwin"
+Requires-Dist: requests (>=2.31.0,<3.0.0)
+Requires-Dist: sentence-transformers (>=2.2.2,<3.0.0)
+Requires-Dist: tokenizers (>=0.13.3,<0.14.0)
+Requires-Dist: tqdm (>=4.65.0,<5.0.0)
+Project-URL: Repository, https://github.com/qdrant/fastembed
+Description-Content-Type: text/markdown
+
+# FastEmbed Library
+
+FastEmbed is a Python library that provides convenient methods for indexing and searching text documents using Qdrant, a high-dimensional vector indexing and search system.
+
+## Features
+
+- Batch document insertion with automatic embedding using SentenceTransformers. With support for OpenAI and custom embeddings.
+- Efficient batch searching with support for filtering by metadata.
+- Automatic generation of unique IDs for documents.
+- Convenient alias methods for adding documents and performing queries.
+
+## Installation
+
+To install the FastEmbed library, we install Qdrant client as well with pip:
+
+```bash
+pip install fastembed qdrant-client
+```
+
+## Usage
+
+Here's a simple usage example, which works as is:
+
+```python
+from qdrant_client import QdrantClient
+
+# Initialize the client
+client = QdrantClient(":memory:")  # or QdrantClient(path="path/to/db")
+
+# Prepare your documents, metadata, and IDs
+docs = ["Qdrant has Langchain integrations", "Qdrant also has Llama Index integrations"]
+metadatas = [
+    {"source": "Langchain-docs"},
+    {"source": "Linkedin-docs"},
+]
+ids = [42, 2]
+
+# Use the new add method
+client.add(collection_name="demo_collection", docs={"documents": docs, "metadatas": metadatas, "ids": ids})
+
+search_result = client.query(collection_name="demo_collection", query_texts=["This is a query document"])
+print(search_result)
+```
+

fastembed-0.0.1/README.md

@@ -0,0 +1,43 @@
+# FastEmbed Library
+
+FastEmbed is a Python library that provides convenient methods for indexing and searching text documents using Qdrant, a high-dimensional vector indexing and search system.
+
+## Features
+
+- Batch document insertion with automatic embedding using SentenceTransformers. With support for OpenAI and custom embeddings.
+- Efficient batch searching with support for filtering by metadata.
+- Automatic generation of unique IDs for documents.
+- Convenient alias methods for adding documents and performing queries.
+
+## Installation
+
+To install the FastEmbed library, we install Qdrant client as well with pip:
+
+```bash
+pip install fastembed qdrant-client
+```
+
+## Usage
+
+Here's a simple usage example, which works as is:
+
+```python
+from qdrant_client import QdrantClient
+
+# Initialize the client
+client = QdrantClient(":memory:")  # or QdrantClient(path="path/to/db")
+
+# Prepare your documents, metadata, and IDs
+docs = ["Qdrant has Langchain integrations", "Qdrant also has Llama Index integrations"]
+metadatas = [
+    {"source": "Langchain-docs"},
+    {"source": "Linkedin-docs"},
+]
+ids = [42, 2]
+
+# Use the new add method
+client.add(collection_name="demo_collection", docs={"documents": docs, "metadatas": metadatas, "ids": ids})
+
+search_result = client.query(collection_name="demo_collection", query_texts=["This is a query document"])
+print(search_result)
+```

fastembed-0.0.1/fastembed/embedding.py

@@ -0,0 +1,207 @@
+import os
+import shutil
+import tarfile
+import tempfile
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any, List
+
+import numpy as np
+import onnxruntime as ort
+import requests
+from tokenizers import Tokenizer
+from tqdm import tqdm
+
+
+# Use pytorches default epsilon for division by zero
+# https://pytorch.org/docs/stable/generated/torch.nn.functional.normalize.html
+def normalize(v):
+    norm = np.linalg.norm(v, axis=1)
+    norm[norm == 0] = 1e-12
+    return v / norm[:, np.newaxis]
+
+class Embedding(ABC):
+    @abstractmethod
+    def encode(self, texts: List[str])->List[np.ndarray]:
+        pass
+
+class DefaultEmbedding(Embedding):
+    def download_file_from_gcs(self, url: str, output_path: str, show_progress: bool=True):
+        if os.path.exists(output_path):
+            return output_path
+        response = requests.get(url, stream=True)
+
+        # Handle HTTP errors
+        if response.status_code == 403:
+            print("Authentication error: you do not have permission to access this resource.")
+            return
+        elif response.status_code != 200:
+            print(f"HTTP error {response.status_code} while trying to download the file.")
+            return
+
+        # Get the total size of the file
+        total_size_in_bytes = int(response.headers.get("content-length", 0))
+
+        # Warn if the total size is zero
+        if total_size_in_bytes == 0:
+            print(f"Warning: Content-length header is missing or zero in the response from {url}.")
+
+        # Initialize the progress bar
+        progress_bar = (
+            tqdm(total=total_size_in_bytes, unit="iB", unit_scale=True) if total_size_in_bytes and show_progress else None
+        )
+
+        # Attempt to download the file
+        try:
+            with open(output_path, "wb") as file:
+                for chunk in response.iter_content(chunk_size=1024):  # Adjust chunk size to your preference
+                    if chunk:  # Filter out keep-alive new chunks
+                        if progress_bar is not None:
+                            progress_bar.update(len(chunk))
+                        file.write(chunk)
+        except Exception as e:
+            print(f"An error occurred while trying to download the file: {str(e)}")
+            return
+        finally:
+            if progress_bar is not None:
+                progress_bar.close()
+        return output_path
+
+    def decompress_to_cache(self, targz_path: str, cache_dir: str=None):
+        # Check if targz_path exists and is a file
+        if not os.path.isfile(targz_path):
+            raise ValueError(f"{targz_path} does not exist or is not a file.")
+
+        # Check if targz_path is a .tar.gz file
+        if not targz_path.endswith(".tar.gz"):
+            raise ValueError(f"{targz_path} is not a .tar.gz file.")
+
+        # Create a temporary directory for caching if cache_dir is not provided
+        if cache_dir is None:
+            cache_dir = tempfile.mkdtemp()
+
+        try:
+            # Open the tar.gz file
+            with tarfile.open(targz_path, "r:gz") as tar:
+                # Extract all files into the cache directory
+                tar.extractall(path=cache_dir)
+            print(f"Files have been decompressed into {cache_dir}")
+        except tarfile.TarError as e:
+            # If any error occurs while opening or extracting the tar.gz file,
+            # delete the cache directory (if it was created in this function)
+            # and raise the error again
+            if "tmp" in cache_dir:
+                shutil.rmtree(cache_dir)
+            raise ValueError(f"An error occurred while decompressing {targz_path}: {e}")
+
+        return cache_dir
+
+    def __init__(self, model_name="sentence-transformers/all-MiniLM-L6-v2"):        
+        filepath = self.download_file_from_gcs(
+            "https://storage.googleapis.com/qdrant-fastembed/fast-all-MiniLM-L6-v2.tar.gz", output_path="fast-all-MiniLM-L6-v2.tar.gz"
+        )
+
+        model_dir = self.decompress_to_cache(targz_path=filepath)
+        
+        model_dir = Path(model_dir) / "fast-all-MiniLM-L6-v2"
+        tokenizer_path = model_dir / "tokenizer.json"
+        if not tokenizer_path.exists():
+            raise ValueError(f"Could not find tokenizer.json in {model_dir}")
+        model_path = model_dir / "model_optimized.onnx"
+        if not model_path.exists():
+            raise ValueError(f"Could not find model_optimized.onnx in {model_dir}")
+
+        self.tokenizer = Tokenizer.from_file(str(tokenizer_path))
+        self.tokenizer.enable_truncation(max_length=256)
+        self.tokenizer.enable_padding(pad_id=0, pad_token="[PAD]", length=256)
+        self.model = ort.InferenceSession(str(model_path), providers=['CoreMLExecutionProvider', 'CPUExecutionProvider'])
+
+    def encode(self, documents: List[str], batch_size: int = 32) -> List[np.ndarray]:
+        # Encode documents using the model
+        # https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2
+        
+        all_embeddings = []
+        for i in range(0, len(documents), batch_size):
+            batch = documents[i:i + batch_size]
+            encoded = [self.tokenizer.encode(d) for d in batch]
+            input_ids = np.array([e.ids for e in encoded])
+            attention_mask = np.array([e.attention_mask for e in encoded])
+            onnx_input = {
+                "input_ids": np.array(input_ids, dtype=np.int64),
+                "attention_mask": np.array(attention_mask, dtype=np.int64),
+                "token_type_ids": np.array([np.zeros(len(e), dtype=np.int64) for e in input_ids], dtype=np.int64),
+            }
+            model_output = self.model.run(None, onnx_input)
+            last_hidden_state = model_output[0]
+            # Perform mean pooling with attention weighting
+            input_mask_expanded = np.broadcast_to(np.expand_dims(attention_mask, -1), last_hidden_state.shape)
+            embeddings = np.sum(last_hidden_state * input_mask_expanded, 1) / np.clip(input_mask_expanded.sum(1), a_min=1e-9, a_max=None)
+            embeddings = normalize(embeddings).astype(np.float32)
+            all_embeddings.append(embeddings)
+        return np.concatenate(all_embeddings)
+
+
+
+
+
+class SentenceTransformersEmbedding(Embedding):
+    def __init__(self, model_name="sentence-transformers/all-MiniLM-L6-v2"):
+        try:
+            from sentence_transformers import SentenceTransformer
+        except ImportError:
+            raise ImportError("Please install the sentence-transformers package to use this method.")
+        self.model = SentenceTransformer(model_name)
+
+    def encode(self, texts):
+        return self.model.encode(texts)
+
+class GeneralTextEmbedding(Embedding):
+    """
+    https://huggingface.co/thenlper/gte-large
+
+    SoTA embedding model for text based retrieval tasks.
+    """
+
+    @classmethod
+    def average_pool(last_hidden_states: Any,
+                    attention_mask: Any) -> Any:
+        last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
+        return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
+
+    def __init__(self, model_name="thenlper/gte-large"):
+        try:
+            import torch.nn.functional as F
+            from transformers import AutoModel, AutoTokenizer
+        except ImportError:
+            raise ImportError("Please install the transformers package with torch to use this method.")
+        self.tokenizer = AutoTokenizer.from_pretrained("thenlper/gte-large")
+        self.model = AutoModel.from_pretrained("thenlper/gte-large")
+        
+    def encode(self, input_texts: List[str]):
+        try:
+            import torch.nn.functional as F
+        except ImportError:
+            raise ImportError("Please install Pytorch to use this method.")
+        # Tokenize the input texts
+        batch_dict = self.tokenizer(input_texts, max_length=512, padding=True, truncation=True, return_tensors='pt')
+
+        outputs = self.model(**batch_dict)
+        embeddings = GeneralTextEmbedding.average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
+
+        # (Optionally) normalize embeddings
+        embeddings = F.normalize(embeddings, p=2, dim=1)
+        scores = (embeddings[:1] @ embeddings[1:].T) * 100
+        return scores
+
+        
+
+class OpenAIEmbedding(Embedding):
+    def __init__(self):
+        # Initialize your OpenAI model here
+        # self.model = ...
+        ...
+
+    def encode(self, texts):
+        # Use your OpenAI model to encode the texts
+        # return self.model.encode(texts)
+        raise NotImplementedError

fastembed-0.0.1/fastembed/qdrant_client.py

@@ -0,0 +1,20 @@
+# qdrant_client.py
+
+
+class QdrantClient:
+    # Existing QdrantClient implementation...
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # Check if fastvector is installed
+        try:
+            from fastembed.qdrant_mixin import QdrantClientMixin
+
+            # If it is, add the mixin methods to this instance
+            for name, method in QdrantClientMixin.__dict__.items():
+                if callable(method):
+                    setattr(self, name, method.__get__(self, self.__class__))
+        except ImportError:
+            # If it's not, do nothing
+            pass

fastembed-0.0.1/fastembed/qdrant_mixin.py

@@ -0,0 +1,150 @@
+import uuid
+from typing import Any, Dict, Generator, List, Optional
+
+from pydantic import BaseModel
+from qdrant_client.models import Distance, PointStruct, SearchParams, VectorParams
+
+from .embedding import DefaultEmbedding, Embedding
+
+
+class QueryResponse(BaseModel):
+    ids: List[str]
+    embeddings: Optional[List[List[float]]]
+    metadatas: List[Dict[str, Any]]
+    distances: List[float]
+
+
+class QdrantClientMixin:
+    def batch_iterable(self, iterable: List[Any], batch_size: int) -> Generator[List[Any], None, None]:
+        """A generator that yields batches of items from an iterable."""
+        batch = []
+        for item in iterable:
+            batch.append(item)
+            if len(batch) >= batch_size:
+                yield batch
+                batch = []
+        if batch:
+            yield batch
+
+    def create_collection(
+        self,
+        collection_name: str,
+        vectors_config: VectorParams = VectorParams(size=1536, distance=Distance.COSINE),
+        **kwargs,
+    ) -> None:
+        """
+        Create a collection with the given name and vectors config.
+
+        Args:
+            collection_name (str): _description_
+            vectors_config (VectorParams, optional): Defaults to VectorParams(size=1536, distance=Distance.COSINE).
+        """
+        self.recreate_collection(collection_name=collection_name, vectors_config=vectors_config, **kwargs)
+
+    def upsert_docs(
+        self,
+        collection_name: str,
+        docs: Dict[str, List[Any]],
+        batch_size: int = 512,
+        wait: bool = True,
+        embedding_model: Embedding = DefaultEmbedding(),
+        **kwargs,
+    ) -> None:
+        """
+        Args:
+            collection_name (str): _description_
+            docs (Dict[str, List[Any]]): _description_
+            batch_size (int, optional): _description_. Defaults to 512.
+            wait (bool, optional): _description_. Defaults to True.
+            embedding_model (Embedding, optional): Defaults to DefaultEmbedding() with all-MiniLM-L6-v2.
+        """
+
+        # Iterate over documents and metadatas in batches
+        for batch_docs, batch_metadatas in zip(
+            self.batch_iterable(docs["documents"], batch_size),
+            self.batch_iterable(docs["metadatas"], batch_size),
+        ):
+            # Tokenize, embed, and index each document
+            embeddings = embedding_model.encode(batch_docs)
+
+            # Create a PointStruct for each document
+            points = [
+                PointStruct(id=uuid.uuid4().hex, vector=embedding.tolist(), payload={**metadata, "text": doc})
+                for doc, embedding, metadata in zip(batch_docs, embeddings, batch_metadatas)
+            ]
+
+            # Check if collection exists
+            if collection_name not in self.get_collections():
+                self.create_collection(
+                    collection_name=collection_name,
+                    vectors_config=VectorParams(size=len(embeddings[0]), distance=Distance.COSINE),
+                )
+            # Call the existing upsert method with the new PointStruct
+            self.upsert(collection_name=collection_name, points=points, wait=wait, **kwargs)
+
+    def search_docs(
+        self,
+        collection_name: str,
+        query_texts: List[str],
+        n_results: int = 2,
+        batch_size: int = 512,
+        query_filter: Optional[Dict[str, Any]] = None,
+        search_params: SearchParams = SearchParams(hnsw_ef=128, exact=False),
+        embedding_model: Embedding = DefaultEmbedding(),
+        **kwargs,
+    ) -> List[QueryResponse]:
+        """
+        Search for documents in a collection.
+
+        Args:
+            collection_name (str): _description_
+            query_texts (List[str]): _description_
+            n_results (int, optional): _description_. Defaults to 2.
+            query_filter (Optional[Dict[str, Any]], optional): _description_. Defaults to None.
+            search_params (SearchParams, optional): _description_. Defaults to SearchParams(hnsw_ef=128, exact=False).
+            embedding_model (Embedding, optional): _description_. Defaults to SentenceTransformersEmbedding().
+            batch_size (int, optional): _description_. Defaults to 512.
+
+        Returns:
+            List[QueryResponse]: _description_
+        """
+
+        query_responses = []
+
+        # Perform the search for each batch of query texts
+        for query_texts_batch in self.batch_iterable(query_texts, batch_size):
+            query_vectors = embedding_model.encode(query_texts_batch)
+
+            for _, query_vector in zip(query_texts_batch, query_vectors):
+                search_result = self.search(
+                    collection_name=collection_name,
+                    query_filter=query_filter,
+                    search_params=search_params,
+                    query_vector=query_vector.tolist(),
+                    limit=n_results,
+                    with_payload=True,
+                    **kwargs,
+                )
+
+                ids, embeddings, metadatas, distances = [], [], [], []
+                for scored_point in search_result:
+                    ids.append(scored_point.id)
+                    if scored_point.vector:
+                        embeddings.append(scored_point.vector)
+                    metadatas.append(scored_point.payload)
+                    distances.append(scored_point.score)
+
+                query_responses.append(
+                    QueryResponse(
+                        ids=ids,
+                        embeddings=embeddings,
+                        metadatas=metadatas,
+                        distances=distances,
+                    )
+                )
+
+        return query_responses
+
+    # Define aliases for the methods
+    add = upsert_docs
+    query = search_docs

fastembed-0.0.1/pyproject.toml

@@ -0,0 +1,42 @@
+[tool.poetry]
+name = "fastembed"
+version = "0.0.1"
+description = "Fast Embedding Creation and Simpler API for Qdrant"
+authors = ["NirantK <nirant.bits@gmail.com>"]
+license = "Apache License"
+readme = "README.md"
+packages = [{include = "fastembed"}]
+homepage = "https://github.com/qdrant/fastembed"
+repository = "https://github.com/qdrant/fastembed"
+keywords = ["vector", "embedding", "neural", "search", "qdrant", "sentence-transformers"]
+
+[tool.poetry.dependencies]
+python = ">=3.8.0,<3.12"
+onnxruntime = "^1.15.1"
+sentence-transformers = "^2.2.2"
+tqdm = "^4.65.0"
+requests = "^2.31.0"
+tokenizers = "^0.13.3"
+
+[tool.poetry.dev-dependencies]
+ruff = "^0.0.277"
+isort = "^5.12.0"
+black = "^23.7.0"
+pre-commit = "^3.3.3"
+
+[tool.poetry.dependencies.onnxruntime-silicon]
+version = "^1.15.0"
+markers = "sys_platform == 'darwin'"  # This makes it macOS specific
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.black]
+line-length = 120
+
+[tool.isort]
+profile = "black"
+
+[tool.ruff]
+line-length = 120