feat: add fastembed embeddings (#374)

Author: mhordynski

.libraries-whitelist.txt

@@ -5,3 +5,4 @@ chroma-hnswlib
 rouge
 distilabel
 rerankers
+py_rust_stemmers

packages/ragbits-core/CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Add support to fastembed dense & sparse embeddings.
 - Fix: changed variable type from Filter to WhereQuery in the Qdrant vector store in list method.
 
 ## 0.8.0 (2025-01-29)

packages/ragbits-core/pyproject.toml

@@ -54,6 +54,9 @@ local = [
     "transformers~=4.44.2",
     "numpy~=1.26.0"
 ]
+fastembed = [
+    "fastembed>=0.4.2"
+]
 lab = [
     "gradio~=4.44.0",
 ]

packages/ragbits-core/src/ragbits/core/embeddings/fastembed.py

@@ -0,0 +1,77 @@
+from fastembed import SparseTextEmbedding, TextEmbedding
+
+from ragbits.core.embeddings import Embeddings, EmbeddingsOptionsT, SparseEmbeddings
+from ragbits.core.embeddings.sparse import SparseVector
+from ragbits.core.options import Options
+
+
+class FastEmbedOptions(Options):
+    """
+    Dataclass that represents available call options for the LocalEmbeddings client.
+    """
+
+    batch_size: int = 256
+    parallel: int | None = None
+
+
+class FastEmbedEmbeddings(Embeddings[FastEmbedOptions]):
+    """
+    Class for creating dense text embeddings using FastEmbed library.
+    For more information, see the [FastEmbed GitHub](https://github.com/qdrant/fastembed).
+    """
+
+    options_cls = FastEmbedOptions
+    _model: TextEmbedding
+
+    def __init__(self, model_name: str, default_options: FastEmbedOptions | None = None):
+        super().__init__(default_options=default_options)
+        self.model_name = model_name
+        self._model = TextEmbedding(model_name)
+
+    async def embed_text(self, data: list[str], options: EmbeddingsOptionsT | None = None) -> list[list[float]]:
+        """
+        Embeds a list of strings into a list of embeddings.
+
+        Args:
+            data: List of strings to get embeddings for.
+            options: Additional options to pass to the embedding model.
+
+        Returns:
+            List of embeddings for the given strings.
+        """
+        merged_options = (self.default_options | options) if options else self.default_options
+
+        return [[float(x) for x in result] for result in self._model.embed(data, **merged_options.dict())]
+
+
+class FastEmbedSparseEmbeddings(SparseEmbeddings[FastEmbedOptions]):
+    """
+    Class for creating sparse text embeddings using FastEmbed library.
+    For more information, see the [FastEmbed GitHub](https://github.com/qdrant/fastembed).
+    """
+
+    options_cls = FastEmbedOptions
+    _model: SparseTextEmbedding
+
+    def __init__(self, model_name: str, default_options: FastEmbedOptions | None = None):
+        super().__init__(default_options=default_options)
+        self.model_name = model_name
+        self._model = SparseTextEmbedding(model_name)
+
+    async def embed_text(self, data: list[str], options: EmbeddingsOptionsT | None = None) -> list[SparseVector]:
+        """
+        Embeds a list of strings into a list of sparse embeddings.
+
+        Args:
+            data: List of strings to get embeddings for.
+            options: Additional options to pass to the embedding model.
+
+        Returns:
+            List of embeddings for the given strings.
+        """
+        merged_options = (self.default_options | options) if options else self.default_options
+
+        return [
+            SparseVector(values=[float(x) for x in result.values], indices=[int(x) for x in result.indices])
+            for result in self._model.embed(data, **merged_options.dict())
+        ]

packages/ragbits-core/src/ragbits/core/embeddings/sparse.py

@@ -1,9 +1,9 @@
 from abc import ABC, abstractmethod
 from collections import Counter
-from dataclasses import dataclass
 from typing import ClassVar, TypeVar
 
 import tiktoken
+from pydantic import BaseModel
 
 from ragbits.core import embeddings
 from ragbits.core.options import Options
@@ -13,22 +13,18 @@
 SparseEmbeddingsOptionsT = TypeVar("SparseEmbeddingsOptionsT", bound=Options)
 
 
-@dataclass
-class SparseVector:
+class SparseVector(BaseModel):
     """Sparse Vector representation"""
 
-    non_zero_dims: list[int]
-    non_zero_vals: list[int]
-    dim: int
+    indices: list[int]
+    values: list[float]
 
     def __post_init__(self) -> None:
-        if len(self.non_zero_dims) != len(self.non_zero_vals):
+        if len(self.indices) != len(self.values):
             raise ValueError("There should be the same number of non-zero values as non-zero positions")
-        if any(dim >= self.dim or dim < 0 for dim in self.non_zero_dims):
-            raise ValueError("Indexes should be in the range of the vector dim")
 
     def __repr__(self) -> str:
-        return f"SparseVector(non_zero_dims={self.non_zero_dims}, non_zero_vals={self.non_zero_vals}, dim={self.dim})"
+        return f"SparseVector(indices={self.indices}, values={self.values})"
 
 
 class SparseEmbeddings(ConfigurableComponent[SparseEmbeddingsOptionsT], ABC):
@@ -39,7 +35,7 @@ class SparseEmbeddings(ConfigurableComponent[SparseEmbeddingsOptionsT], ABC):
     configuration_key: ClassVar = "sparse_embedder"
 
     @abstractmethod
-    def embed_text(self, texts: list[str], options: SparseEmbeddingsOptionsT | None = None) -> list[SparseVector]:
+    async def embed_text(self, texts: list[str], options: SparseEmbeddingsOptionsT | None = None) -> list[SparseVector]:
         """Transforms a list of texts into sparse vectors"""
 
 
@@ -52,11 +48,11 @@ class BagOfTokensOptions(Options):
 
 
 class BagOfTokens(SparseEmbeddings[BagOfTokensOptions]):
-    """BagofTokens implementations of sparse Embeddings interface"""
+    """BagOfTokens implementations of sparse Embeddings interface"""
 
     options_cls = BagOfTokensOptions
 
-    def embed_text(self, texts: list[str], options: BagOfTokensOptions | None = None) -> list[SparseVector]:
+    async def embed_text(self, texts: list[str], options: BagOfTokensOptions | None = None) -> list[SparseVector]:
         """
         Transforms a list of texts into sparse vectors using bag-of-tokens representation.
 
@@ -73,12 +69,14 @@ def embed_text(self, texts: list[str], options: BagOfTokensOptions | None = None
             raise ValueError("Please specify only one of encoding_name or model_name")
         if not (merged_options.encoding_name or merged_options.model_name):
             raise ValueError("Either encoding_name or model_name needs to be specified")
+
         if merged_options.encoding_name:
             encoder = tiktoken.get_encoding(encoding_name=merged_options.encoding_name)
-        if merged_options.model_name:
+        elif merged_options.model_name:
             encoder = tiktoken.encoding_for_model(model_name=merged_options.model_name)
+        else:
+            raise ValueError("Either encoding_name or model_name needs to be specified")
 
-        dim = encoder.n_vocab
         min_token_count = merged_options.min_token_count or float("-inf")
         for text in texts:
             tokens = encoder.encode(text)
@@ -90,7 +88,7 @@ def embed_text(self, texts: list[str], options: BagOfTokensOptions | None = None
                 if count < min_token_count:
                     continue
                 non_zero_dims.append(token)
-                non_zero_vals.append(count)
+                non_zero_vals.append(float(count))
 
-            vectors.append(SparseVector(non_zero_dims, non_zero_vals, dim))
+            vectors.append(SparseVector(indices=non_zero_dims, values=non_zero_vals))
         return vectors

packages/ragbits-core/tests/unit/embeddings/__init__.py

@@ -0,0 +1,13 @@
+from ragbits.core.embeddings.fastembed import FastEmbedEmbeddings, FastEmbedSparseEmbeddings
+
+
+async def test_fastembed_dense_embeddings():
+    embeddings = FastEmbedEmbeddings("BAAI/bge-small-en-v1.5")
+    result = await embeddings.embed_text(["text1"])
+    assert len(result[0]) == 384
+
+
+async def test_fastembed_sparse_embeddings():
+    embeddings = FastEmbedSparseEmbeddings("qdrant/bm25")
+    result = await embeddings.embed_text(["text1"])
+    assert len(result[0].values) == len(result[0].indices)

packages/ragbits-core/tests/unit/embeddings/fastembed.py

@@ -7,11 +7,11 @@ requires-python = ">=3.10"
 dependencies = [
     "asyncpg>=0.30.0",
     "ragbits-cli",
-    "ragbits-core[chroma,lab,local,otel,qdrant]",
+    "ragbits-core[chroma,lab,fastembed,local,otel,qdrant]",
     "ragbits-document-search[gcs,huggingface,distributed,azure,s3]",
     "ragbits-evaluate[relari]",
     "ragbits-guardrails[openai]",
-    "ragbits-conversations"
+    "ragbits-conversations",
 ]
 
 [tool.uv]