added basic test

koaning · Apr 5, 2024 · d92bdae · d92bdae
1 parent 5a29e15
commit d92bdae
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 14 deletions.
diff --git a/embetter/text/__init__.py b/embetter/text/__init__.py
@@ -1,9 +1,5 @@
 from embetter.error import NotInstalled
-
-try:
-    from embetter.text._sbert import SentenceEncoder
-except ModuleNotFoundError:
-    SentenceEncoder = NotInstalled("SentenceEncoder", "sentence-tfm")
+from embetter.text._sbert import SentenceEncoder, MatrouskaEncoder
 
 try:
     from embetter.text._s2v import Sense2VecEncoder
@@ -36,6 +32,7 @@
 
 __all__ = [
     "SentenceEncoder",
+    "MatrouskaEncoder",
     "Sense2VecEncoder",
     "BytePairEncoder",
     "spaCyEncoder",

diff --git a/embetter/text/_sbert.py b/embetter/text/_sbert.py
@@ -7,10 +7,6 @@
 from embetter.base import EmbetterBase
 
 
-class MatrouskaEncoder(EmbetterBase):
-    def __init__(self, name="tomaarsen/mpnet-base-nli-matryoshka", **kwargs):
-        return SentenceEncoder(name=name, **kwargs)
-
 class SentenceEncoder(EmbetterBase):
     """
     Encoder that can numerically encode sentences.
@@ -98,3 +94,7 @@ def transform(self, X, y=None):
             X = X.to_numpy()
 
         return self.tfm.encode(X)
+
+
+def MatrouskaEncoder(name="tomaarsen/mpnet-base-nli-matryoshka", **kwargs):
+    return SentenceEncoder(name=name, **kwargs)
diff --git a/tests/test_text.py b/tests/test_text.py
@@ -13,6 +13,7 @@
     SentenceEncoder,
     GensimEncoder,
     spaCyEncoder,
+    MatrouskaEncoder,
     learn_lite_text_embeddings,
     LiteTextEncoder,
 )
@@ -46,17 +47,18 @@ def test_word2vec(setting):
     assert repr(encoder)
 
 
-def test_basic_sentence_encoder():
+@pytest.mark.parametrize("encoder", [MatrouskaEncoder, SentenceEncoder])
+def test_basic_sentence_encoder(encoder):
     """Check correct dimensions and repr for SentenceEncoder."""
-    encoder = SentenceEncoder()
+    enc = encoder()
     # Embedding dim of underlying model
-    output_dim = encoder.tfm._modules["1"].word_embedding_dimension
-    output = encoder.fit_transform(test_sentences)
+    output_dim = enc.tfm._modules["1"].word_embedding_dimension
+    output = enc.fit_transform(test_sentences)
     assert isinstance(output, np.ndarray)
     assert output.shape == (len(test_sentences), output_dim)
     # scikit-learn configures repr dynamically from defined attributes.
     # To test correct implementation we should test if calling repr breaks.
-    assert repr(encoder)
+    assert repr(enc)
 
 
 @pytest.mark.parametrize("setting", ["max", "mean", "both"])