add experimental tqdm

kartynnik · Dec 11, 2018 · 182db68 · 182db68
1 parent 2bd4ff5
commit 182db68
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/setup_google_colab.py b/setup_google_colab.py
@@ -14,7 +14,10 @@ def download_github_code(path):
 
 
 def setup_common():
-    os.system("pip install tqdm")
+    if bool(int(os.environ.get("EXPERIMENTAL_TQDM", "0"))):
+        os.system("pip install --force https://github.com/hse-aml/intro-to-dl/releases/download/ColabTqdm/tqdm-colab.zip")
+    else:
+        os.system("pip install tqdm")
     os.system("pip install --upgrade Keras==2.0.6")  # latest version breaks callbacks
     download_github_code("keras_utils.py")
     download_github_code("grading.py")

diff --git a/tqdm_utils.py b/tqdm_utils.py
@@ -47,16 +47,17 @@ def __next__(self):
             raise StopIteration
 
 
-def is_colab():
+def use_simple_tqdm():
     try:
         import google.colab
-        return True
+        import os
+        return not bool(int(os.environ.get("EXPERIMENTAL_TQDM", "0")))
     except ImportError:
         return False
 
 
 def tqdm_notebook_failsafe(*args, **kwargs):
-    if is_colab():
+    if use_simple_tqdm():
         # tqdm is broken on Google Colab
         return SimpleTqdm(*args, **kwargs)
     else: