Adds PreTrainedModel.framework attribute (#13817)

* Added `framework` attribute * Update modeling_utils.py * Update modeling_flax_utils.py * Update modeling_tf_utils.py * Update modeling_utils.py * Update modeling_tf_utils.py * Update modeling_tf_utils.py * Update modeling_flax_utils.py * Update modeling_tf_utils.py * Update modeling_utils.py * Update modeling_utils.py * Update modeling_tf_utils.py * Update modeling_flax_utils.py * string -> str * Update modeling_tf_utils.py * string -> str * fixup * make flake happy Co-authored-by: patil-suraj <surajp815@gmail.com>
huggingface · Oct 8, 2021 · de34481 · de34481
1 parent d70919e
commit de34481
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 0 deletions.
diff --git a/src/transformers/modeling_flax_utils.py b/src/transformers/modeling_flax_utils.py
@@ -119,6 +119,13 @@ def _from_config(cls, config, **kwargs):
         """
         return cls(config, **kwargs)
 
+    @property
+    def framework(self) -> str:
+        """
+        :str: Identifies that this is a Flax model.
+        """
+        return "flax"
+
     @property
     def config(self) -> PretrainedConfig:
         return self._config

diff --git a/src/transformers/modeling_tf_utils.py b/src/transformers/modeling_tf_utils.py
@@ -652,6 +652,13 @@ def dummy_inputs(self) -> Dict[str, tf.Tensor]:
             "input_ids": tf.constant(DUMMY_INPUTS),
         }
 
+    @property
+    def framework(self) -> str:
+        """
+        :str: Identifies that this is a TensorFlow model.
+        """
+        return "tf"
+
     def __init__(self, config, *inputs, **kwargs):
         super().__init__(*inputs, **kwargs)
         if not isinstance(config, PretrainedConfig):

diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
@@ -461,6 +461,13 @@ def dummy_inputs(self) -> Dict[str, torch.Tensor]:
         """
         return {"input_ids": torch.tensor(DUMMY_INPUTS)}
 
+    @property
+    def framework(self) -> str:
+        """
+        :str: Identifies that this is a PyTorch model.
+        """
+        return "pt"
+
     def __init__(self, config: PretrainedConfig, *inputs, **kwargs):
         super().__init__()
         if not isinstance(config, PretrainedConfig):