refactor: move BagOfTokens model_name / encoding_name parameters to init (#592)

Author: mhordynski

packages/ragbits-core/CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Refacor: move BagOfTokens model_name / encoding_name parameters to init (#592)
 - Update utils (#590)
 - Resolve vector_size by PgVectorStore automatically (#588)
 - Add get_vector_size method to all Embedders (#587)

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

@@ -12,8 +12,6 @@
 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,31 +20,48 @@ class BagOfTokens(SparseEmbedder[BagOfTokensOptions]):
 
     options_cls = BagOfTokensOptions
 
-    async def get_vector_size(self) -> VectorSize:
+    def __init__(
+        self,
+        model_name: str | None = None,
+        encoding_name: str | None = None,
+        default_options: BagOfTokensOptions | None = None,
+    ) -> None:
         """
-        Get the vector size for this BagOfTokens model.
+        Initialize the BagOfTokens embedder.
 
-        For BagOfTokens, this returns the tokenizer vocabulary size.
+        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.
 
-        Returns:
-            VectorSize object with is_sparse=True and the vocabulary size.
+        Raises:
+            ValueError: If both model_name and encoding_name are provided, or if neither is provided.
         """
-        merged_options = self.default_options
+        super().__init__(default_options=default_options)
 
-        if merged_options.encoding_name and merged_options.model_name:
+        if encoding_name and 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)
+        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")
 
-        # Get the vocabulary size from the encoder
-        vocab_size = encoder.n_vocab
+    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]:
@@ -63,21 +78,9 @@ async def embed_text(self, texts: list[str], options: BagOfTokensOptions | None
         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 = []

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

@@ -2,13 +2,11 @@
 
 from ragbits.core.embeddings.base import VectorSize
 from ragbits.core.embeddings.sparse.bag_of_tokens import BagOfTokens, BagOfTokensOptions
-from ragbits.core.types import NOT_GIVEN
 
 
 async def test_bag_of_tokens_get_vector_size_with_encoding():
     """Test BagOfTokens get_vector_size method with encoding_name."""
-    options = BagOfTokensOptions(encoding_name="cl100k_base", model_name=NOT_GIVEN)
-    embedder = BagOfTokens(default_options=options)
+    embedder = BagOfTokens(encoding_name="cl100k_base")
 
     vector_size = await embedder.get_vector_size()
 
@@ -20,8 +18,7 @@ async def test_bag_of_tokens_get_vector_size_with_encoding():
 
 async def test_bag_of_tokens_get_vector_size_with_model():
     """Test BagOfTokens get_vector_size method with model_name."""
-    options = BagOfTokensOptions(model_name="gpt-3.5-turbo")
-    embedder = BagOfTokens(default_options=options)
+    embedder = BagOfTokens(model_name="gpt-3.5-turbo")
 
     vector_size = await embedder.get_vector_size()
 
@@ -44,26 +41,22 @@ async def test_bag_of_tokens_get_vector_size_default():
 
 async def test_bag_of_tokens_get_vector_size_error_both_specified():
     """Test BagOfTokens get_vector_size raises error when both encoding_name and model_name are specified."""
-    options = BagOfTokensOptions(encoding_name="cl100k_base", model_name="gpt-3.5-turbo")
-    embedder = BagOfTokens(default_options=options)
-
     with pytest.raises(ValueError, match="Please specify only one of encoding_name or model_name"):
-        await embedder.get_vector_size()
+        BagOfTokens(encoding_name="cl100k_base", model_name="gpt-3.5-turbo")
 
 
 async def test_bag_of_tokens_get_vector_size_error_none_specified():
     """Test BagOfTokens get_vector_size raises error when neither encoding_name nor model_name are specified."""
-    options = BagOfTokensOptions(encoding_name=NOT_GIVEN, model_name=NOT_GIVEN)
-    embedder = BagOfTokens(default_options=options)
-
-    with pytest.raises(ValueError, match="Either encoding_name or model_name needs to be specified"):
-        await embedder.get_vector_size()
+    # This test is no longer valid since we now default to gpt-4o when nothing is specified
+    # The constructor will automatically use gpt-4o as default
+    embedder = BagOfTokens()
+    vector_size = await embedder.get_vector_size()
+    assert vector_size.size > 0  # Should succeed with default gpt-4o
 
 
 async def test_bag_of_tokens_embed_text_consistency():
     """Test that BagOfTokens embeddings are consistent with vector size."""
-    options = BagOfTokensOptions(encoding_name="cl100k_base", model_name=NOT_GIVEN)
-    embedder = BagOfTokens(default_options=options)
+    embedder = BagOfTokens(encoding_name="cl100k_base")
 
     # Get vector size
     vector_size = await embedder.get_vector_size()
@@ -79,15 +72,26 @@ async def test_bag_of_tokens_embed_text_consistency():
 
 async def test_bag_of_tokens_different_encodings():
     """Test BagOfTokens with different encodings have different vocabulary sizes."""
-    options1 = BagOfTokensOptions(encoding_name="cl100k_base", model_name=NOT_GIVEN)
-    embedder1 = BagOfTokens(default_options=options1)
-
-    options2 = BagOfTokensOptions(encoding_name="p50k_base", model_name=NOT_GIVEN)
-    embedder2 = BagOfTokens(default_options=options2)
+    embedder1 = BagOfTokens(encoding_name="cl100k_base")
+    embedder2 = BagOfTokens(encoding_name="p50k_base")
 
     vector_size1 = await embedder1.get_vector_size()
     vector_size2 = await embedder2.get_vector_size()
 
     assert vector_size1.size != vector_size2.size
     assert vector_size1.is_sparse is True
     assert vector_size2.is_sparse is True
+
+
+async def test_bag_of_tokens_min_token_count_option():
+    """Test BagOfTokens with min_token_count option."""
+    embedder = BagOfTokens(encoding_name="cl100k_base")
+    options = BagOfTokensOptions(min_token_count=2)
+
+    # Test with text that has some repeated tokens
+    embeddings = await embedder.embed_text(["test test test"], options=options)
+
+    # Should have embeddings (non-empty vectors)
+    assert len(embeddings) == 1
+    assert len(embeddings[0].indices) > 0
+    assert len(embeddings[0].values) > 0

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

@@ -1,4 +1,5 @@
 import litellm
+import pytest
 
 from ragbits.core.embeddings import DenseEmbedder, NoopEmbedder
 from ragbits.core.embeddings.dense import LiteLLMEmbedder, LiteLLMEmbedderOptions
@@ -44,6 +45,7 @@ def test_subclass_from_config_bag_of_tokens():
         {
             "type": "ragbits.core.embeddings.sparse:BagOfTokens",
             "config": {
+                "model_name": "gpt-4o",
                 "default_options": {
                     "option1": "value1",
                     "option2": "value2",
@@ -54,14 +56,30 @@ def test_subclass_from_config_bag_of_tokens():
     embedder: SparseEmbedder = SparseEmbedder.subclass_from_config(config)
     assert isinstance(embedder, BagOfTokens)
     assert embedder.default_options == BagOfTokensOptions(
-        model_name="gpt-4o",
-        encoding_name=NOT_GIVEN,
         min_token_count=NOT_GIVEN,
         option1="value1",
         option2="value2",
     )  # type: ignore
 
 
+def test_subclass_from_config_bag_of_tokens_both_specified():
+    config = ObjectConstructionConfig.model_validate(
+        {
+            "type": "ragbits.core.embeddings.sparse:BagOfTokens",
+            "config": {
+                "model_name": "gpt-4o",
+                "encoding_name": "cl100k_base",
+                "default_options": {
+                    "option1": "value1",
+                    "option2": "value2",
+                },
+            },
+        }
+    )
+    with pytest.raises(ValueError, match="Please specify only one of encoding_name or model_name"):
+        SparseEmbedder.subclass_from_config(config)
+
+
 def test_from_config_with_router():
     config = ObjectConstructionConfig(
         type="ragbits.core.embeddings.dense:LiteLLMEmbedder",