downloader: better progress bar

This commit is contained in:
oobabooga 2024-07-28 22:21:56 -07:00
parent 9dcff21da9
commit f4d95f33b8

View File

@ -212,11 +212,15 @@ class ModelDownloader:
total_size = int(r.headers.get('content-length', 0)) total_size = int(r.headers.get('content-length', 0))
block_size = 1024 * 1024 # 1MB block_size = 1024 * 1024 # 1MB
filename_str = str(filename) # Convert PosixPath to string if necessary
tqdm_kwargs = { tqdm_kwargs = {
'total': total_size, 'total': total_size,
'unit': 'iB', 'unit': 'B',
'unit_scale': True, 'unit_scale': True,
'bar_format': '{l_bar}{bar}| {n_fmt}/{total_fmt} {rate_fmt}' 'unit_divisor': 1024,
'bar_format': '{desc}{percentage:3.0f}%|{bar:50}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]',
'desc': f"{filename_str}: "
} }
if 'COLAB_GPU' in os.environ: if 'COLAB_GPU' in os.environ:
@ -233,7 +237,7 @@ class ModelDownloader:
t.update(len(data)) t.update(len(data))
if total_size != 0 and self.progress_bar is not None: if total_size != 0 and self.progress_bar is not None:
count += len(data) count += len(data)
self.progress_bar(float(count) / float(total_size), f"{filename}") self.progress_bar(float(count) / float(total_size), f"{filename_str}")
break # Exit loop if successful break # Exit loop if successful
except (RequestException, ConnectionError, Timeout) as e: except (RequestException, ConnectionError, Timeout) as e: