Add from_dict tests

deepset-ai · Aug 2, 2024 · 923b112 · 923b112
1 parent bacde0f
commit 923b112
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/test/components/embedders/test_sentence_transformers_document_embedder.py b/test/components/embedders/test_sentence_transformers_document_embedder.py
@@ -132,6 +132,8 @@ def test_from_dict(self):
             "meta_fields_to_embed": ["meta_field"],
             "trust_remote_code": True,
             "truncate_dim": 256,
+            "model_kwargs": {"torch_dtype": "torch.float32"},
+            "tokenizer_kwargs": {"model_max_length": 512},
         }
         component = SentenceTransformersDocumentEmbedder.from_dict(
             {
@@ -151,6 +153,8 @@ def test_from_dict(self):
         assert component.trust_remote_code
         assert component.meta_fields_to_embed == ["meta_field"]
         assert component.truncate_dim == 256
+        assert component.model_kwargs == {"torch_dtype": torch.float32}
+        assert component.tokenizer_kwargs == {"model_max_length": 512}
 
     def test_from_dict_no_default_parameters(self):
         component = SentenceTransformersDocumentEmbedder.from_dict(

diff --git a/test/components/rankers/test_transformers_similarity.py b/test/components/rankers/test_transformers_similarity.py
@@ -156,6 +156,7 @@ def test_from_dict(self):
                 "calibration_factor": None,
                 "score_threshold": 0.01,
                 "model_kwargs": {"torch_dtype": "torch.float16"},
+                "tokenizer_kwargs": {"model_max_length": 512},
             },
         }
 
@@ -176,6 +177,7 @@ def test_from_dict(self):
             "torch_dtype": torch.float16,
             "device_map": ComponentDevice.resolve_device(None).to_hf(),
         }
+        assert component.tokenizer_kwargs == {"model_max_length": 512}
 
     def test_from_dict_no_default_parameters(self):
         data = {