some pr issues (#514)

* some pr issues

* revert query embed refactor

* test: add query embed tests

* nit

* Update tests/test_sparse_embeddings.py

---------

Co-authored-by: Andrey Vasnetsov <andrey@vasnetsov.com>

Author: joein

fastembed/late_interaction/colbert.py

@@ -244,9 +244,7 @@ def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterab
             self.load_onnx_model()
 
         for text in query:
-            yield from self._post_process_onnx_output(
-                self.onnx_embed([text], is_doc=False), is_doc=False
-            )
+            yield from self._post_process_onnx_output(self.onnx_embed([text]), is_doc=False)
 
     @classmethod
     def _get_worker_class(cls) -> Type[TextEmbeddingWorker[NumpyArray]]:

fastembed/late_interaction/token_embeddings.py

@@ -1,5 +1,7 @@
-from typing import Union, Iterable, Optional, List, Dict, Any, Type
+from dataclasses import asdict
+from typing import Union, Iterable, Optional, Any, Type
 
+from fastembed.common.model_description import DenseModelDescription, ModelSource
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.types import NumpyArray
 from fastembed.late_interaction.late_interaction_embedding_base import (
@@ -10,29 +12,38 @@
 import numpy as np
 
 supported_token_embeddings_models = [
-    {
-        "model": "jinaai/jina-embeddings-v2-small-en-tokens",
-        "dim": 512,
-        "description": "Text embeddings, Unimodal (text), English, 8192 input tokens truncation,"
+    DenseModelDescription(
+        model="jinaai/jina-embeddings-v2-small-en-tokens",
+        dim=512,
+        description="Text embeddings, Unimodal (text), English, 8192 input tokens truncation,"
         " Prefixes for queries/documents: not necessary, 2023 year.",
-        "license": "apache-2.0",
-        "size_in_GB": 0.12,
-        "sources": {"hf": "xenova/jina-embeddings-v2-small-en"},
-        "model_file": "onnx/model.onnx",
-    },
+        license="apache-2.0",
+        size_in_GB=0.12,
+        sources=ModelSource(hf="xenova/jina-embeddings-v2-small-en"),
+        model_file="onnx/model.onnx",
+    ),
 ]
 
 
 class TokenEmbeddingsModel(OnnxTextEmbedding, LateInteractionTextEmbeddingBase):
     @classmethod
-    def list_supported_models(cls) -> List[Dict[str, Any]]:
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
         """Lists the supported models.
 
         Returns:
-            List[Dict[str, Any]]: A list of dictionaries containing the model information.
+            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
         """
         return supported_token_embeddings_models
 
+    @classmethod
+    def list_supported_models(cls) -> list[dict[str, Any]]:
+        """Lists the supported models.
+
+        Returns:
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
+        """
+        return [asdict(model) for model in cls._list_supported_models()]
+
     @classmethod
     def _get_worker_class(cls) -> Type[TextEmbeddingWorker[NumpyArray]]:
         return TokensEmbeddingWorker
@@ -47,7 +58,6 @@ def _post_process_onnx_output(
         masks = output.attention_mask
 
         # For each document we only select those embeddings that are not masked out
-
         for i in range(embeddings.shape[0]):
             yield embeddings[i, masks[i] == 1]
 
@@ -58,11 +68,9 @@ def embed(
         parallel: Optional[int] = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
-        yield from OnnxTextEmbedding.embed(
-            self, documents, batch_size=batch_size, parallel=parallel, **kwargs
-        )
+        yield from super().embed(documents, batch_size=batch_size, parallel=parallel, **kwargs)
 
-    def tokenize_docs(self, documents: List[str]) -> List[NumpyArray]:
+    def tokenize_docs(self, documents: list[str]) -> list[NumpyArray]:
         if self.tokenizer is None:
             raise ValueError("Tokenizer not initialized")
         encoded = self.tokenizer.encode_batch(documents)
@@ -83,6 +91,7 @@ def init_embedding(
 
 if __name__ == "__main__":
     # Example usage
+    print(TokenEmbeddingsModel.list_supported_models())
     model = TokenEmbeddingsModel(model_name="jinaai/jina-embeddings-v2-small-en-tokens")
     docs = ["Hello, world!", "hello", "hello hello"]
 

fastembed/sparse/minicoil.py

@@ -1,24 +1,24 @@
 from pathlib import Path
+
+from typing import Any, Optional, Sequence, Iterable, Union, Type
+
 import numpy as np
-from typing import Any, Dict, Optional, Sequence, Iterable, Union, Set
+from numpy.typing import NDArray
+from py_rust_stemmers import SnowballStemmer
+from tokenizers import Tokenizer
 
 from fastembed.common.model_description import SparseModelDescription, ModelSource
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common import OnnxProvider
+from fastembed.common.utils import define_cache_dir
 from fastembed.sparse.sparse_embedding_base import (
     SparseEmbedding,
     SparseTextEmbeddingBase,
 )
-
-from numpy.typing import NDArray
-
-from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.sparse.utils.minicoil_encoder import Encoder
 from fastembed.sparse.utils.sparse_vectors_converter import SparseVectorConverter, WordEmbedding
 from fastembed.sparse.utils.vocab_resolver import VocabResolver, VocabTokenizerTokenizer
 from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
-from py_rust_stemmers import SnowballStemmer
-from fastembed.common import OnnxProvider
-from fastembed.common.utils import define_cache_dir
-from tokenizers import Tokenizer
 
 
 MINICOIL_MODEL_FILE = "minicoil.triplet.model.npy"
@@ -29,7 +29,7 @@
 supported_minicoil_models: list[SparseModelDescription] = [
     SparseModelDescription(
         model="Qdrant/minicoil-v1",
-        vocab_size=30522,
+        vocab_size=19125,
         description="Sparse embedding model, that resolves semantic meaning of the words, "
         "while keeping exact keyword match behavior. "
         "Based on jinaai/jina-embeddings-v2-small-en-tokens",
@@ -57,7 +57,7 @@ class MiniCOIL(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         while keeping exact keyword match behavior.
 
         Each vocabulary token is converted into 4d component of a sparse vector, which is then weighted by the token frequency in the corpus.
-        If the token is not found in the corpus, it is trearted exactly like in BM25.
+        If the token is not found in the corpus, it is treated exactly like in BM25.
     `
         The model is based on `jinaai/jina-embeddings-v2-small-en-tokens`
     """
@@ -116,10 +116,10 @@ def __init__(
 
         # Initialize class attributes
         self.tokenizer: Optional[Tokenizer] = None
-        self.invert_vocab: Dict[int, str] = {}
-        self.special_tokens: Set[str] = set()
-        self.special_tokens_ids: Set[int] = set()
-        self.stopwords: Set[str] = set()
+        self.invert_vocab: dict[int, str] = {}
+        self.special_tokens: set[str] = set()
+        self.special_tokens_ids: set[int] = set()
+        self.stopwords: set[str] = set()
         self.vocab_resolver: Optional[VocabResolver] = None
         self.encoder: Optional[Encoder] = None
         self.output_dim: Optional[int] = None
@@ -297,7 +297,7 @@ def _post_process_onnx_output(
             # Size of counts: (unique_words)
             words_ids = ids_mapping[:, 0].tolist()
 
-            sentence_result: Dict[str, WordEmbedding] = {}
+            sentence_result: dict[str, WordEmbedding] = {}
 
             words = [self.vocab_resolver.lookup_word(word_id) for word_id in words_ids]
 
@@ -325,36 +325,25 @@ def _post_process_onnx_output(
                     word=oov_word, forms=[oov_word], count=int(count), word_id=-1, embedding=[1]
                 )
 
-            if is_query:
-                yield self.sparse_vector_converter.embedding_to_vector_query(
+            if not is_query:
+                yield self.sparse_vector_converter.embedding_to_vector(
                     sentence_result, vocab_size=vocab_size, embedding_size=embedding_size
                 )
             else:
-                yield self.sparse_vector_converter.embedding_to_vector(
+                yield self.sparse_vector_converter.embedding_to_vector_query(
                     sentence_result, vocab_size=vocab_size, embedding_size=embedding_size
                 )
 
+    @classmethod
+    def _get_worker_class(cls) -> Type["MiniCoilTextEmbeddingWorker"]:
+        return MiniCoilTextEmbeddingWorker
+
 
 class MiniCoilTextEmbeddingWorker(TextEmbeddingWorker[SparseEmbedding]):
     def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> MiniCOIL:
         return MiniCOIL(
             model_name=model_name,
             cache_dir=cache_dir,
+            threads=1,
             **kwargs,
         )
-
-
-def test_minicoil() -> None:
-    model = MiniCOIL(model_name="Qdrant/minicoil-v1")
-
-    embedding = next(iter(model.embed("Hello World")))
-
-    print(embedding)
-
-    embedding = next(iter(model.query_embed("Hello World")))
-
-    print(embedding)
-
-
-if __name__ == "__main__":
-    test_minicoil()

fastembed/sparse/utils/minicoil_encoder.py

@@ -7,8 +7,6 @@
 import numpy as np
 from fastembed.common.types import NumpyArray
 
-from typing import Tuple
-
 
 class Encoder:
     """
@@ -35,7 +33,7 @@ class Encoder:
      │                     │
      └─────────────────────┘
 
-     Final liner transformation is accompanied by a non-linear activation function: Tanh.
+     Final linear transformation is accompanied by a non-linear activation function: Tanh.
 
      Tanh is used to ensure that the output is in the range [-1, 1].
      It would be easier to visually interpret the output of the model, assuming that each dimension
@@ -70,7 +68,7 @@ def convert_vocab_ids(vocab_ids: NumpyArray) -> NumpyArray:
     @classmethod
     def avg_by_vocab_ids(
         cls, vocab_ids: NumpyArray, embeddings: NumpyArray
-    ) -> Tuple[NumpyArray, NumpyArray]:
+    ) -> tuple[NumpyArray, NumpyArray]:
         """
         Takes:
             vocab_ids: (batch_size, seq_len) int array
@@ -112,7 +110,7 @@ def avg_by_vocab_ids(
 
     def forward(
         self, vocab_ids: NumpyArray, embeddings: NumpyArray
-    ) -> Tuple[NumpyArray, NumpyArray]:
+    ) -> tuple[NumpyArray, NumpyArray]:
         """
         Args:
             vocab_ids: (batch_size, seq_len) int array

fastembed/sparse/utils/sparse_vectors_converter.py

@@ -32,7 +32,7 @@ def __init__(
         avg_len: float = 150.0,
     ):
         punctuation = set(get_all_punctuation())
-        special_tokens = set(["[CLS]", "[SEP]", "[PAD]", "[UNK]", "[MASK]"])
+        special_tokens = {"[CLS]", "[SEP]", "[PAD]", "[UNK]", "[MASK]"}
 
         self.stemmer = stemmer
         self.unwanted_tokens = punctuation | special_tokens | stopwords
@@ -163,20 +163,21 @@ def embedding_to_vector(
         unknown_words_shift = (
             (vocab_size * embedding_size) // GAP + 2
         ) * GAP  # miniCOIL vocab + at least (GAP // embedding_size) + 1 new words gap
-
         sentence_embedding_cleaned = self.clean_words(sentence_embedding)
 
-        # Calcualte sentence length after cleaning
+        # Calculate sentence length after cleaning
         sentence_len = 0
         for embedding in sentence_embedding_cleaned.values():
             sentence_len += embedding.count
 
         for embedding in sentence_embedding_cleaned.values():
             word_id = embedding.word_id
-            num_occurences = embedding.count
-            tf = self.bm25_tf(num_occurences, sentence_len)
-
-            if word_id >= 0:  # miniCOIL starts with ID 1
+            num_occurrences = embedding.count
+            tf = self.bm25_tf(num_occurrences, sentence_len)
+            if (
+                word_id > 0
+            ):  # miniCOIL starts with ID 1, we generally won't have word_id == 0 (UNK), as we don't add
+                # these words to sentence_embedding
                 embedding_values = embedding.embedding
                 normalized_embedding = self.normalize_vector(embedding_values)