feat: add kwargs to SentenceTransformerEmbeddingModel init

NVIDIA · Jan 9, 2025 · db6245a · db6245a
1 parent a6116c1
commit db6245a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/nemoguardrails/embeddings/providers/sentence_transformers.py b/nemoguardrails/embeddings/providers/sentence_transformers.py
@@ -41,7 +41,7 @@ class SentenceTransformerEmbeddingModel(EmbeddingModel):
 
     engine_name = "SentenceTransformers"
 
-    def __init__(self, embedding_model: str):
+    def __init__(self, embedding_model: str, **kwargs):
         try:
             from sentence_transformers import SentenceTransformer
         except ImportError:
@@ -58,7 +58,7 @@ def __init__(self, embedding_model: str):
             )
 
         device = "cuda" if cuda.is_available() else "cpu"
-        self.model = SentenceTransformer(embedding_model, device=device)
+        self.model = SentenceTransformer(embedding_model, device=device, **kwargs)
         # Get the embedding dimension of the model
         self.embedding_size = self.model.get_sentence_embedding_dimension()