From 072669eeda149a8a1aa582e2dab381306d229a05 Mon Sep 17 00:00:00 2001
From: Benjamin Fineran <bfineran@users.noreply.github.com>
Date: Tue, 25 Jan 2022 21:22:49 -0500
Subject: [PATCH] update transformers install flow to latest changes from
 sparseml (#252) (#253)

* update transformers install flow to latest changes from sparseml

* test fixes
---
 src/deepsparse/transformers/__init__.py | 43 ++++++++++++++++---------
 1 file changed, 28 insertions(+), 15 deletions(-)

diff --git a/src/deepsparse/transformers/__init__.py b/src/deepsparse/transformers/__init__.py
index 6d43b2c9ff..a3a7213d1a 100644
--- a/src/deepsparse/transformers/__init__.py
+++ b/src/deepsparse/transformers/__init__.py
@@ -25,35 +25,46 @@
 try:
     import transformers as _transformers
 
+    # triggers error if neuralmagic/transformers is not installed
+    _transformers.models.bert.modeling_bert.QATMatMul
     _transformers_import_error = None
 except Exception as _transformers_import_err:
     _transformers_import_error = _transformers_import_err
 
 
 _LOGGER = _logging.getLogger(__name__)
+_NM_TRANSFORMERS_TAR_TEMPLATE = (
+    "https://github.com/neuralmagic/transformers/releases/download/"
+    "{version}/transformers-4.7.0.dev0-py3-none-any.whl"
+)
+_NM_TRANSFORMERS_NIGHTLY = _NM_TRANSFORMERS_TAR_TEMPLATE.format(version="nightly")
 
 
 def _install_transformers_and_deps():
 
+    import subprocess as _subprocess
+    import sys as _sys
+
     import deepsparse as _deepsparse
-    import pip as _pip
 
-    transformers_branch = (
-        "master"
+    nm_transformers_release = (
+        "nightly"
         if not _deepsparse.is_release
-        else f"release/{_deepsparse.version.version_major_minor}"
+        else f"v{_deepsparse.version.version_major_minor}"
     )
-    transformers_requirement = (
-        "transformers @ git+https://github.com/neuralmagic/transformers.git"
-        f"@{transformers_branch}"
+    transformers_requirement = _NM_TRANSFORMERS_TAR_TEMPLATE.format(
+        version=nm_transformers_release
     )
 
     try:
-        _pip.main(
+        _subprocess.check_call(
             [
+                _sys.executable,
+                "-m",
+                "pip",
                 "install",
                 transformers_requirement,
-                "datasets",
+                "datasets<1.18.0",
                 "sklearn",
                 "seqeval",
             ]
@@ -66,7 +77,7 @@ def _install_transformers_and_deps():
         raise ValueError(
             "Unable to install and import deepsparse-transformers dependencies check "
             "that transformers is installed, if not, install via "
-            "`pip install git+https://github.com/neuralmagic/transformers.git`"
+            f"`pip install {_NM_TRANSFORMERS_NIGHTLY}`"
         )
 
 
@@ -82,13 +93,15 @@ def _check_transformers_install():
             # skip any further checks
             return
         else:
-            _LOGGER.info(
-                "No installation of transformers found. Installing deepsparse-transformers "
-                "dependencies"
+            _LOGGER.warning(
+                "deepsparse-transformers installation not detected. Installing "
+                "deepsparse-transformers dependencies if transformers is already "
+                "installed in the environment, it will be overwritten. Set "
+                "environment variable NM_NO_AUTOINSTALL_TRANSFORMERS to disable"
             )
             _install_transformers_and_deps()
 
-    # check NM fork installed with QATMatMul available
+    # re check import after potential install
     try:
         import transformers as _transformers
 
@@ -98,7 +111,7 @@ def _check_transformers_install():
             "transformers.models.bert.modeling_bert.QATMatMul not availalbe. the"
             "neuralmagic fork of transformers may not be installed. it can be "
             "installed via "
-            "`pip install git+https://github.com/neuralmagic/transformers.git`"
+            f"`pip install {_NM_TRANSFORMERS_NIGHTLY}`"
         )