Allow hidding frameworks through environment variables (NO_TF, NO_TORCH).

transformers/file_utils.py

@@ -27,17 +27,25 @@ from contextlib import contextmanager
 logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
 
 try:
-    import tensorflow as tf
-    assert hasattr(tf, '__version__') and int(tf.__version__[0]) >= 2
-    _tf_available = True  # pylint: disable=invalid-name
-    logger.info("TensorFlow version {} available.".format(tf.__version__))
+    if 'NO_TF' in os.environ and os.environ['NO_TF'].upper() in ('1', 'ON'):
+        logger.info("Found NO_TF, disabling TensorFlow")
+        _tf_available = False
+    else:
+        import tensorflow as tf
+        assert hasattr(tf, '__version__') and int(tf.__version__[0]) >= 2
+        _tf_available = True  # pylint: disable=invalid-name
+        logger.info("TensorFlow version {} available.".format(tf.__version__))
 except (ImportError, AssertionError):
     _tf_available = False  # pylint: disable=invalid-name
 
 try:
-    import torch
-    _torch_available = True  # pylint: disable=invalid-name
-    logger.info("PyTorch version {} available.".format(torch.__version__))
+    if 'NO_TORCH' in os.environ and os.environ['NO_TORCH'].upper() in ('1', 'ON'):
+        logger.info("Found NO_TORCH, disabling PyTorch")
+        _torch_available = False
+    else:
+        import torch
+        _torch_available = True  # pylint: disable=invalid-name
+        logger.info("PyTorch version {} available.".format(torch.__version__))
 except ImportError:
     _torch_available = False  # pylint: disable=invalid-name
 
@@ -77,6 +85,7 @@ def is_torch_available():
     return _torch_available
 
 def is_tf_available():
+
     return _tf_available
 
 if not six.PY2: