from_pretrained flag in TransformersEmbedder, to avoid downloadin…

…g the model if you want to start from a checkpoint.

.gitignore

@@ -3,6 +3,8 @@ stuff
 /test.ipynb
 /test.py
 
+# Fleet
+.fleet
 
 # Created by https://www.toptal.com/developers/gitignore/api/python,pycharm+all,vscode,macos,linux,windows
 # Edit at https://www.toptal.com/developers/gitignore?templates=python,pycharm+all,vscode,macos,linux,windows

requirements.txt

@@ -1,2 +1,2 @@
 torch>=1.7,<1.13
-transformers>=4.3,<4.21
+transformers>=4.14,<4.21

setup.py

@@ -7,11 +7,11 @@
 extras["torch"] = ["torch>=1.5,<1.13"]
 extras["all"] = extras["torch"]
 
-install_requires = ["transformers>=4.3,<4.21"]
+install_requires = ["transformers>=4.14,<4.21"]
 
 setuptools.setup(
     name="transformers_embedder",
-    version="3.0.2",
+    version="3.0.3",
     author="Riccardo Orlando",
     author_email="[email protected]",
     description="Word level transformer based embeddings",

transformers_embedder/embedder.py

@@ -52,6 +52,10 @@ class TransformersEmbedder(torch.nn.Module):
             If ``True``, the transformer model is fine-tuned during training.
         return_all (:obj:`bool`, optional, defaults to :obj:`False`):
             If ``True``, returns all the outputs from the HuggingFace model.
+        from_pretrained (:obj:`bool`, optional, defaults to :obj:`True`):
+            If ``True``, the model is loaded from a pre-trained model, otherwise it is initialized with
+            random weights. Usefull when you want to load a model from a specific checkpoint, without
+            having to download the entire model.
     """
 
     def __init__(
@@ -62,6 +66,7 @@ def __init__(
         output_layers: Sequence[int] = (-4, -3, -2, -1),
         fine_tune: bool = True,
         return_all: bool = False,
+        from_pretrained: bool = True,
         *args,
         **kwargs,
     ) -> None:
@@ -74,9 +79,14 @@ def __init__(
                 *args,
                 **kwargs,
             )
-            self.transformer_model = tr.AutoModel.from_pretrained(
-                model, config=config, *args, **kwargs
-            )
+            if from_pretrained:
+                self.transformer_model = tr.AutoModel.from_pretrained(
+                    model, config=config, *args, **kwargs
+                )
+            else:
+                self.transformer_model = tr.AutoModel.from_config(
+                    config, *args, **kwargs
+                )
         else:
             self.transformer_model = model