Small fix on hyperparameter search (#150)

* edit modeling.py

* add unit test

Co-authored-by: Mouhaned Chebaane <[email protected]>

Author: Mouhanedg56

src/setfit/modeling.py

@@ -376,7 +376,7 @@ def _from_pretrained(
 
                     model_head = multilabel_classifier
                 else:
-                    model_head = LogisticRegression()
+                    model_head = clf
 
         return SetFitModel(model_body=model_body, model_head=model_head, multi_target_strategy=multi_target_strategy)
 

tests/test_modeling.py

@@ -53,6 +53,24 @@ def test_setfit_default_model_head():
     assert type(model.model_head) is LogisticRegression
 
 
+def test_setfit_model_head_params():
+    params = {
+        "head_params": {
+            "max_iter": 200,
+            "solver": "newton-cg",
+        }
+    }
+
+    model = SetFitModel.from_pretrained("sentence-transformers/paraphrase-albert-small-v2", **params)
+
+    assert type(model.model_head) is LogisticRegression
+    assert params["head_params"] == {
+        parameter: value
+        for parameter, value in model.model_head.get_params(deep=False).items()
+        if parameter in params["head_params"]
+    }
+
+
 def test_setfit_multilabel_one_vs_rest_model_head():
     model = SetFitModel.from_pretrained(
         "sentence-transformers/paraphrase-albert-small-v2", multi_target_strategy="one-vs-rest"