From 8661276b20188644f51e3828527bb2e5b6885282 Mon Sep 17 00:00:00 2001 From: chenyu Date: Thu, 5 Sep 2024 21:43:31 -0400 Subject: [PATCH] set tqdm total to None in fetch gunzip (#6377) fixed bar overflow python examples/beautiful_mnist.py https://storage.googleapis.com/cvdf-datasets/mnist/train-images-idx3-ubyte.gz: 47.0MB [00:00, 136MB/s]] https://storage.googleapis.com/cvdf-datasets/mnist/train-labels-idx1-ubyte.gz: 60.0kB [00:00, 8.66MB/s] https://storage.googleapis.com/cvdf-datasets/mnist/t10k-images-idx3-ubyte.gz: 7.84MB [00:00, 97.5MB/s] https://storage.googleapis.com/cvdf-datasets/mnist/t10k-labels-idx1-ubyte.gz: 10.0kB [00:00, 8.90MB/s] --- tinygrad/helpers.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index 4d4261fd3c..da61cd2ec1 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -272,15 +272,15 @@ def fetch(url:str, name:Optional[Union[pathlib.Path, str]]=None, subdir:Optional if not fp.is_file() or not allow_caching: with urllib.request.urlopen(url, timeout=10) as r: assert r.status == 200 - total_length = int(r.headers.get('content-length', 0)) - progress_bar = tqdm(total=total_length, unit='B', unit_scale=True, desc=f"{url}", disable=CI) + length = int(r.headers.get('content-length', 0)) if not gunzip else None + progress_bar = tqdm(total=length, unit='B', unit_scale=True, desc=f"{url}", disable=CI) (path := fp.parent).mkdir(parents=True, exist_ok=True) readfile = gzip.GzipFile(fileobj=r) if gunzip else r with tempfile.NamedTemporaryFile(dir=path, delete=False) as f: while chunk := readfile.read(16384): progress_bar.update(f.write(chunk)) f.close() progress_bar.update(close=True) - if (file_size:=os.stat(f.name).st_size) < total_length: raise RuntimeError(f"fetch size incomplete, {file_size} < {total_length}") + if length and (file_size:=os.stat(f.name).st_size) < length: raise RuntimeError(f"fetch size incomplete, {file_size} < {length}") pathlib.Path(f.name).rename(fp) return fp