Merge pull request #27 from stefan-it/show-download-progress

data-generators: show nice progress bar for download process

Author: lukaszkaiser

tensor2tensor/data_generators/generator_utils.py

@@ -126,6 +126,18 @@ def generate_files(generator,
   return output_files
 
 
+def download_report_hook(count, block_size, total_size):
+  """Report hook for download progress
+
+  Args:
+    count: current block number
+    block_size: block size
+    total_size: total size
+  """
+  percent = int(count*block_size*100/total_size)
+  print("\r%d%%" % percent + ' completed', end='\r')
+
+
 def maybe_download(directory, filename, url):
   """Download filename from url unless it's already in directory.
 
@@ -143,7 +155,11 @@ def maybe_download(directory, filename, url):
   filepath = os.path.join(directory, filename)
   if not tf.gfile.Exists(filepath):
     tf.logging.info("Downloading %s to %s" % (url, filepath))
-    filepath, _ = urllib.urlretrieve(url, filepath)
+    filepath, _ = urllib.urlretrieve(url, filepath,
+                                     reporthook=download_report_hook)
+
+    # Print newline to clear the carriage return from the download progress
+    print()
     statinfo = os.stat(filepath)
     tf.logging.info("Succesfully downloaded %s, %s bytes." % (filename,
                                                               statinfo.st_size))