Mirrors/text-generation-webui
1
0
Fork 0
mirror of https://github.com/oobabooga/text-generation-webui.git synced 2024-11-26 09:40:20 +01:00
Code Issues Packages Projects Releases Wiki Activity

tqdm improvement for colab

Browse Source
This commit is contained in:
oobabooga 2023-10-21 22:00:29 -07:00
parent ae79c510cc
commit cd45635f53
1 changed files with 15 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
download-model.py
View File

@ -167,8 +167,22 @@ class ModelDownloader:
r.raise_for_status() # Do not continue the download if the request was unsuccessful r.raise_for_status() # Do not continue the download if the request was unsuccessful
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
tqdm_kwargs = {
'total': total_size,
'unit': 'iB',
'unit_scale': True,
'bar_format': '{l_bar}{bar}| {n_fmt:6}/{total_fmt:6} {rate_fmt:6}'
}
if 'COLAB_GPU' in os.environ:
tqdm_kwargs.update({
'position': 0,
'leave': True
})
with open(output_path, mode) as f: with open(output_path, mode) as f:
with tqdm.tqdm(total=total_size, unit='iB', unit_scale=True, bar_format='{l_bar}{bar}| {n_fmt:6}/{total_fmt:6} {rate_fmt:6}') as t: with tqdm.tqdm(**tqdm_kwargs) as t:
count = 0 count = 0
for data in r.iter_content(block_size): for data in r.iter_content(block_size):
t.update(len(data)) t.update(len(data))