Add SetFitModelConfig

src/setfit/config.py

@@ -0,0 +1,20 @@
+from dataclasses import dataclass
+from typing import Any, Optional
+
+from transformers import PretrainedConfig
+
+
+@dataclass
+class SetFitConfig:
+    """
+    Configuration for SetFitModel.
+
+    Args:
+        model_body (`PretrainedConfig`):
+            Configuration of the Sentence Transformer body.
+        model_head (`Optional[Any]`):
+            Configuration of the head.
+    """
+
+    model_body: PretrainedConfig
+    model_head: Optional[Any]

src/setfit/modeling.py

@@ -4,6 +4,8 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Dict, List, Optional, Union
 
+from .config import SetFitConfig
+
 
 # Google Colab runs on Python 3.7, so we need this to be compatible
 try:
@@ -211,6 +213,20 @@ def __init__(
         self.model_original_state = copy.deepcopy(self.model_body.state_dict())
         self.normalize_embeddings = normalize_embeddings
 
+    @property
+    def config(self) -> SetFitConfig:
+        model_body_config = self.model_body._modules["0"]._modules["auto_model"].config
+
+        model_head_config = None
+        if isinstance(self.model_head, SetFitHead):
+            model_head_config = self.model_head.get_config_dict()
+
+        return SetFitConfig(model_body=model_body_config, model_head=model_head_config)
+
+    @config.setter
+    def config(self, config: SetFitConfig) -> None:
+        self.model_body._modules["0"]._modules["auto_model"].config = config.model_body
+
     def fit(
         self,
         x_train: List[str],