Expose transformer model config

Riccorl · May 19, 2023 · 77cc2a1 · 77cc2a1
1 parent 8d599d4
commit 77cc2a1
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/setup.py b/setup.py
@@ -12,7 +12,7 @@
 
 setuptools.setup(
     name="transformers_embedder",
-    version="3.0.9",
+    version="3.0.10",
     author="Riccardo Orlando",
     author_email="orlandoricc@gmail.com",
     description="Word level transformer based embeddings",

diff --git a/transformers_embedder/embedder.py b/transformers_embedder/embedder.py
@@ -75,7 +75,7 @@ def __init__(
     ) -> None:
         super().__init__()
         if isinstance(model, str):
-            config = tr.AutoConfig.from_pretrained(
+            self.config = tr.AutoConfig.from_pretrained(
                 model,
                 output_hidden_states=True,
                 output_attentions=True,
@@ -84,11 +84,11 @@ def __init__(
             )
             if from_pretrained:
                 self.transformer_model = tr.AutoModel.from_pretrained(
-                    model, config=config, *args, **kwargs
+                    model, config=self.config, *args, **kwargs
                 )
             else:
                 self.transformer_model = tr.AutoModel.from_config(
-                    config, *args, **kwargs
+                    self.config, *args, **kwargs
                 )
         else:
             self.transformer_model = model