Odysseus v1.0

scripts/hf_download.py

@@ -0,0 +1,182 @@
+#!/usr/bin/env python3
+"""Download HuggingFace models with clean pipe-friendly progress output.
+
+Usage:
+    python3 scripts/hf_download.py <repo_id> [--include "pattern"]
+
+Prints lines like:
+    FILE model.safetensors [########------------] 42% 1.23/2.91GB 156.3MB/s
+    DONE /path/to/cached/model
+"""
+import argparse
+import sys
+import time
+import os
+
+
+_last_print = {}
+
+
+class PipeTqdm:
+    """Minimal tqdm replacement that prints simple progress lines to stdout."""
+
+    def __init__(self, *args, **kwargs):
+        self.iterable = args[0] if args else kwargs.get("iterable")
+        self.total = kwargs.get("total", None)
+        self.desc = kwargs.get("desc", "")
+        self.unit = kwargs.get("unit", "it")
+        self.n = 0
+        self.start_t = time.time()
+        self.disable = False
+        self._closed = False
+
+        if self.iterable is not None and self.total is None:
+            try:
+                self.total = len(self.iterable)
+            except (TypeError, AttributeError):
+                pass
+
+    def __iter__(self):
+        if self.iterable is None:
+            return
+        for item in self.iterable:
+            yield item
+            self.update(1)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.close()
+
+    def __len__(self):
+        return self.total or 0
+
+    def update(self, n=1):
+        self.n += n
+        total = self.total or 0
+        if total == 0:
+            return
+        now = time.time()
+        key = id(self)
+        # Throttle to every 0.5s, always print on completion
+        if now - _last_print.get(key, 0) < 0.5 and self.n < total:
+            return
+        _last_print[key] = now
+
+        pct = int(100 * self.n / total)
+        elapsed = now - self.start_t
+        speed = self.n / elapsed if elapsed > 0 else 0
+        desc = (self.desc or "").strip()
+
+        # Format sizes
+        if total >= 1024 ** 3:
+            done_s = f"{self.n / (1024**3):.2f}"
+            total_s = f"{total / (1024**3):.2f}GB"
+            speed_s = f"{speed / (1024**2):.1f}MB/s"
+        elif total >= 1024 ** 2:
+            done_s = f"{self.n / (1024**2):.1f}"
+            total_s = f"{total / (1024**2):.1f}MB"
+            speed_s = f"{speed / (1024**2):.1f}MB/s"
+        else:
+            done_s = str(self.n)
+            total_s = str(total)
+            speed_s = f"{speed:.0f}/s"
+
+        # ASCII progress bar
+        bar_len = 20
+        filled = int(bar_len * self.n / total)
+        bar = "#" * filled + "-" * (bar_len - filled)
+
+        print(f"FILE {desc} [{bar}] {pct}% {done_s}/{total_s} {speed_s}", flush=True)
+
+    def set_description(self, desc=None, refresh=True):
+        self.desc = desc or ""
+
+    def set_postfix(self, *args, **kwargs):
+        pass
+
+    def set_postfix_str(self, s="", refresh=True):
+        pass
+
+    def reset(self, total=None):
+        self.n = 0
+        if total is not None:
+            self.total = total
+        self.start_t = time.time()
+
+    def refresh(self):
+        pass
+
+    def close(self):
+        self._closed = True
+
+    def clear(self):
+        pass
+
+    def display(self, msg=None, pos=None):
+        pass
+
+    @property
+    def format_dict(self):
+        return {"n": self.n, "total": self.total, "elapsed": time.time() - self.start_t}
+
+
+def _patch_tqdm():
+    """Replace tqdm everywhere with our pipe-friendly version."""
+    import tqdm as tqdm_mod
+
+    # Replace the main class
+    tqdm_mod.tqdm = PipeTqdm
+    tqdm_mod.auto.tqdm = PipeTqdm
+
+    # huggingface_hub uses tqdm.auto or its own utils.tqdm
+    try:
+        import huggingface_hub.utils
+        huggingface_hub.utils.tqdm = PipeTqdm
+        # Also patch the _tqdm module if it exists
+        if hasattr(huggingface_hub.utils, "_tqdm"):
+            huggingface_hub.utils._tqdm.tqdm = PipeTqdm
+    except (ImportError, AttributeError):
+        pass
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("repo_id", help="HuggingFace repo (e.g. meta-llama/Llama-3-8B)")
+    parser.add_argument("--include", help="File pattern to include (e.g. '*Q4_K_M*')")
+    args = parser.parse_args()
+
+    # Disable HF progress bars (we provide our own)
+    os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "0"
+
+    # Enable Rust-backed parallel downloader if available — big throughput win.
+    # Must be set before importing huggingface_hub.
+    try:
+        import hf_transfer  # noqa: F401
+        os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
+    except ImportError:
+        print("HINT pip install hf_transfer for faster downloads", flush=True)
+
+    _patch_tqdm()
+
+    from huggingface_hub import snapshot_download
+
+    kwargs = {
+        "repo_id": args.repo_id,
+        "max_workers": int(os.environ.get("HF_HUB_DOWNLOAD_MAX_WORKERS", "8")),
+    }
+    if args.include:
+        kwargs["allow_patterns"] = [args.include]
+
+    print(f"START {args.repo_id}", flush=True)
+    try:
+        path = snapshot_download(**kwargs)
+        print(f"DONE {path}", flush=True)
+    except Exception as e:
+        print(f"ERROR {e}", file=sys.stderr, flush=True)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()