SentenceTransformerEmbedder.__init__() kwargs pass-through (#27)

Added: SentenceTransformerConfig init_kwargs attribute/parameter
Changed: SentenceTransformerEmbedder passes config.init_kwargs attribute through to SentenceTransformer.__init__()

Author: andrewschenck

pyproject.toml

@@ -32,7 +32,7 @@ exclude = ["tests", "work"]
 
 [project]
 name = "ragl"
-version = "0.10.0"
+version = "0.10.1"
 dependencies = [
     "bleach",
     "numpy",

ragl/config.py

@@ -23,7 +23,7 @@
 
 import logging
 import re
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any
 
 from ragl.exceptions import ConfigurationError
@@ -104,7 +104,8 @@ class SentenceTransformerConfig(EmbedderConfig):
             to a model on disk.
 
         cache_maxsize:
-            Maximum number of embeddings to cache in memory.
+            Maximum number of entries to cache in memory. Set to
+            0 to disable caching.
         device:
             Device to use for embedding.
         auto_clear_cache:
@@ -116,14 +117,18 @@ class SentenceTransformerConfig(EmbedderConfig):
             Threshold for memory usage before cleaning up cache.
             This is a float between 0.0 and 1.0, where 1.0 means
             100% memory usage.
+       init_kwargs:
+            Additional keyword arguments to pass to the
+            SentenceTransformer constructor.
     """
 
     model_name_or_path: str = 'all-mpnet-base-v2'
-    cache_maxsize: int = 10_000  # set this to 0 to disable caching
+    cache_maxsize: int = 10_000
     device: str | None = None
     auto_clear_cache: bool = True
     show_progress: bool = False
     memory_threshold: float = 0.9
+    init_kwargs: dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self) -> None:
         """Validate configuration after initialization."""

ragl/embed/sentencetransformer.py

@@ -91,7 +91,9 @@ class SentenceTransformerEmbedder:
     def dimensions(self) -> int:
         """Retrieve the embedding dimension count."""
         dimensions = self.model.get_sentence_embedding_dimension()
-        assert isinstance(dimensions, int)
+        if not isinstance(dimensions, int) or dimensions <= 0:
+            raise ValueError('Invalid embedding dimensions '
+                             'retrieved from model')
         return dimensions
 
     def __init__(self, config: SentenceTransformerConfig) -> None:
@@ -111,11 +113,20 @@ def __init__(self, config: SentenceTransformerConfig) -> None:
             _LOG.info('Cache disabled (maxsize=%d)', config.cache_maxsize)
 
         model_path = Path(config.model_name_or_path)
-        self.model = SentenceTransformer(str(model_path), device=config.device)
+        kwargs_device = config.init_kwargs.pop('device', None)
+        if kwargs_device is not None:
+            _LOG.warning('Ignoring device setting in init_kwargs (%s); '
+                         'use config.device (%s) instead',
+                         kwargs_device, config.device)
+        self.model = SentenceTransformer(
+            model_name_or_path=str(model_path),
+            device=config.device,
+            **config.init_kwargs,
+        )
         self._cache_size = config.cache_maxsize
-        self._memory_threshold = config.memory_threshold
         self._auto_cleanup = config.auto_clear_cache
         self._show_progress = config.show_progress
+        self._memory_threshold = config.memory_threshold
         self._embed_cached = lru_cache(self._cache_size)(self._embed_impl)
         _LOG.debug('Embedder initialized: dims=%d, cache_size=%d, device=%s',
                    self.dimensions, self._cache_size, config.device)

tests/functional/ragl/embed/test_sentencetransformer.py

@@ -52,7 +52,7 @@ def test_init(self, mock_sentence_transformer):
         embedder = SentenceTransformerEmbedder(self.config)
 
         # Verify model initialization
-        mock_sentence_transformer.assert_called_once_with("all-MiniLM-L6-v2",
+        mock_sentence_transformer.assert_called_once_with(model_name_or_path="all-MiniLM-L6-v2",
                                                           device="cpu")
         self.assertEqual(embedder.model, self.mock_model)
         self.assertEqual(embedder._cache_size, 100)
@@ -74,7 +74,7 @@ def test_init_with_path(self, mock_sentence_transformer):
 
         embedder = SentenceTransformerEmbedder(config)
 
-        mock_sentence_transformer.assert_called_once_with("all-MiniLM-L6-v2",
+        mock_sentence_transformer.assert_called_once_with(model_name_or_path="all-MiniLM-L6-v2",
                                                           device="cuda")
         self.assertEqual(embedder._cache_size, 50)
         self.assertEqual(embedder._memory_threshold, 0.9)
@@ -102,9 +102,41 @@ def test_init_cache_disabled_logging(self, mock_log,
 
         # Verify embedder was still initialized properly
         self.assertEqual(embedder._cache_size, 0)
-        mock_sentence_transformer.assert_called_once_with("all-MiniLM-L6-v2",
+        mock_sentence_transformer.assert_called_once_with(model_name_or_path="all-MiniLM-L6-v2",
                                                           device="cpu")
 
+    @patch('ragl.embed.sentencetransformer.SentenceTransformer')
+    @patch('ragl.embed.sentencetransformer._LOG')
+    def test_init_kwargs_device_warning(self, mock_log,
+                                        mock_sentence_transformer):
+        """Test that device in init_kwargs triggers warning and is ignored."""
+        mock_sentence_transformer.return_value = self.mock_model
+
+        config = SentenceTransformerConfig(
+            model_name_or_path='all-MiniLM-L6-v2',
+            cache_maxsize=100,
+            memory_threshold=0.8,
+            auto_clear_cache=True,
+            device="cpu",
+            init_kwargs={"device": "cuda", "trust_remote_code": True}
+        )
+
+        embedder = SentenceTransformerEmbedder(config)
+
+        # Verify warning was logged
+        mock_log.warning.assert_called_once_with(
+            'Ignoring device setting in init_kwargs (%s); '
+            'use config.device (%s) instead',
+            'cuda', 'cpu'
+        )
+
+        # Verify SentenceTransformer was called with config.device, not init_kwargs device
+        mock_sentence_transformer.assert_called_once_with(
+            model_name_or_path="all-MiniLM-L6-v2",
+            device="cpu",  # Should use config.device
+            trust_remote_code=True  # Other init_kwargs should still be passed
+        )
+
     @patch('ragl.embed.sentencetransformer.SentenceTransformer')
     def test_dimensions_property(self, mock_sentence_transformer):
         """Test dimensions property."""
@@ -123,7 +155,7 @@ def test_dimensions_property_assertion(self, mock_sentence_transformer):
         self.mock_model.get_sentence_embedding_dimension.return_value = "invalid"
 
         # Now the assertion should fail during initialization
-        with self.assertRaises(AssertionError):
+        with self.assertRaises(ValueError):
             embedder = SentenceTransformerEmbedder(self.config)
 
     @patch('ragl.embed.sentencetransformer.SentenceTransformer')