ragbits-core 0.16.0__py3-none-any.whl → 1.4.0.dev202512021005__py3-none-any.whl

ragbits/core/embeddings/sparse/bag_of_tokens.py

@@ -2,8 +2,8 @@ from collections import Counter
 
 import tiktoken
 
-from ragbits.core.audit import trace
-from ragbits.core.embeddings.base import SparseVector
+from ragbits.core.audit.traces import trace
+from ragbits.core.embeddings.base import SparseVector, VectorSize
 from ragbits.core.embeddings.sparse.base import SparseEmbedder
 from ragbits.core.options import Options
 from ragbits.core.types import NOT_GIVEN, NotGiven
@@ -12,8 +12,6 @@ from ragbits.core.types import NOT_GIVEN, NotGiven
 class BagOfTokensOptions(Options):
     """A dataclass with definition of BOT options"""
 
-    model_name: str | None | NotGiven = "gpt-4o"
-    encoding_name: str | None | NotGiven = NOT_GIVEN
     min_token_count: int | None | NotGiven = NOT_GIVEN
 
 
@@ -22,6 +20,50 @@ class BagOfTokens(SparseEmbedder[BagOfTokensOptions]):
 
     options_cls = BagOfTokensOptions
 
+    def __init__(
+        self,
+        model_name: str | None = None,
+        encoding_name: str | None = None,
+        default_options: BagOfTokensOptions | None = None,
+    ) -> None:
+        """
+        Initialize the BagOfTokens embedder.
+
+        Args:
+            model_name: Name of the model to use for tokenization (e.g., "gpt-4o").
+            encoding_name: Name of the encoding to use for tokenization.
+            default_options: Default options for the embedder.
+
+        Raises:
+            ValueError: If both model_name and encoding_name are provided, or if neither is provided.
+        """
+        super().__init__(default_options=default_options)
+
+        if encoding_name and model_name:
+            raise ValueError("Please specify only one of encoding_name or model_name")
+        if not (encoding_name or model_name):
+            # Default to gpt-4o if neither is specified
+            model_name = "gpt-4o"
+
+        if encoding_name:
+            self._encoder = tiktoken.get_encoding(encoding_name=encoding_name)
+        elif model_name:
+            self._encoder = tiktoken.encoding_for_model(model_name=model_name)
+        else:
+            raise ValueError("Either encoding_name or model_name needs to be specified")
+
+    async def get_vector_size(self) -> VectorSize:
+        """
+        Get the vector size for this BagOfTokens model.
+
+        For BagOfTokens, this returns the tokenizer vocabulary size.
+
+        Returns:
+            VectorSize object with is_sparse=True and the vocabulary size.
+        """
+        vocab_size = self._encoder.n_vocab
+        return VectorSize(size=vocab_size, is_sparse=True)
+
     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.
@@ -36,21 +78,9 @@ class BagOfTokens(SparseEmbedder[BagOfTokensOptions]):
         vectors = []
         merged_options = self.default_options | options if options else self.default_options
         with trace(data=texts, options=merged_options.dict()) as outputs:
-            if merged_options.encoding_name and merged_options.model_name:
-                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)
-            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")
-
             min_token_count = merged_options.min_token_count or float("-inf")
             for text in texts:
-                tokens = encoder.encode(text)
+                tokens = self._encoder.encode(text)
                 token_counts = Counter(tokens)
                 non_zero_dims = []
                 non_zero_vals = []

ragbits/core/embeddings/sparse/base.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 from typing import TypeVar
 
-from ragbits.core.embeddings.base import Embedder, SparseVector
+from ragbits.core.embeddings.base import Embedder, SparseVector, VectorSize
 from ragbits.core.options import Options
 
 SparseEmbedderOptionsT = TypeVar("SparseEmbedderOptionsT", bound=Options)
@@ -23,6 +23,15 @@ class SparseEmbedder(Embedder[SparseEmbedderOptionsT], ABC):
             list of sparse embeddings.
         """
 
+    @abstractmethod
+    async def get_vector_size(self) -> VectorSize:
+        """
+        Get information about the sparse vector size/dimensions returned by this embedder.
+
+        Returns:
+            VectorSize object with is_sparse=True and the vocabulary size.
+        """
+
     async def embed_image(
         self, images: list[bytes], options: SparseEmbedderOptionsT | None = None
     ) -> list[SparseVector]:

ragbits/core/embeddings/sparse/fastembed.py

@@ -2,8 +2,8 @@ from collections.abc import Callable
 
 from fastembed import SparseTextEmbedding
 
-from ragbits.core.audit import trace
-from ragbits.core.embeddings.base import EmbedderOptionsT, SparseVector
+from ragbits.core.audit.traces import trace
+from ragbits.core.embeddings.base import EmbedderOptionsT, SparseVector, VectorSize
 from ragbits.core.embeddings.dense.fastembed import FastEmbedOptions
 from ragbits.core.embeddings.sparse.base import SparseEmbedder
 
@@ -35,6 +35,29 @@ class FastEmbedSparseEmbedder(SparseEmbedder[FastEmbedOptions]):
         """
         return (self.__class__, (self.model_name, self.use_gpu, self.default_options))
 
+    async def get_vector_size(self) -> VectorSize:
+        """
+        Get the vector size for this FastEmbed sparse model.
+
+        For sparse models, this returns the vocabulary size.
+
+        Returns:
+            VectorSize object with is_sparse=True and the vocabulary size.
+        """
+        # Get model info from FastEmbed's supported models list
+        supported_models = self._model.list_supported_models()
+        model_info = next((model for model in supported_models if model["model"] == self.model_name), None)
+
+        if model_info and "vocab_size" in model_info:
+            vocab_size = model_info["vocab_size"]
+        else:
+            sample_embedding = await self.embed_text(["sample text with various tokens"])
+            vocab_size = (
+                max(sample_embedding[0].indices) + 1 if sample_embedding and sample_embedding[0].indices else 30000
+            )
+
+        return VectorSize(size=vocab_size, is_sparse=True)
+
     async def embed_text(self, data: list[str], options: EmbedderOptionsT | None = None) -> list[SparseVector]:
         """
         Embeds a list of strings into a list of sparse embeddings.

ragbits/core/llms/__init__.py

@@ -1,5 +1,5 @@
-from .base import LLM
+from .base import LLM, ToolCall, Usage
 from .litellm import LiteLLM, LiteLLMOptions
-from .local import LocalLLMOptions
+from .local import LocalLLM, LocalLLMOptions
 
-__all__ = ["LLM", "LiteLLM", "LiteLLMOptions", "LocalLLMOptions"]
+__all__ = ["LLM", "LiteLLM", "LiteLLMOptions", "LocalLLM", "LocalLLMOptions", "ToolCall", "Usage"]