Merge pull request #51 from vcvyg/issue-35-cli-download-progress

feat(setup): show live model download progress

Author: ch-liuzhide

src/hebb/cli/commands/setup.py

@@ -7,6 +7,15 @@
 import click
 from click.core import ParameterSource
 from rich.console import Console
+from rich.progress import (
+    BarColumn,
+    DownloadColumn,
+    Progress,
+    TaskProgressColumn,
+    TextColumn,
+    TimeRemainingColumn,
+    TransferSpeedColumn,
+)
 
 from hebb.config.init import default_init_target, initialize_workspace
 from hebb.config.loader import find_config_file, load_settings, update_config_field
@@ -71,8 +80,11 @@ def setup_cmd(ctx: click.Context, language: str, region: str, profile: str) -> N
         if workspace_model_available(workspace, model_id):
             console.print(f"[green]Model already present:[/] {model_cache_dir(workspace, model_id)}")
         else:
-            console.print(f"[cyan]Downloading embedding model[/] ({_download_tier_hint(model_id)})...")
-            model_path = prefetch_model(model_id, workspace, hf_endpoint=region_selection.hf_endpoint)
+            model_path = _prefetch_with_progress(
+                model_id,
+                workspace,
+                hf_endpoint=region_selection.hf_endpoint,
+            )
             console.print(f"[green]Model ready:[/] {model_path}")
         dimension = _verify_model(model_id, hf_endpoint=region_selection.hf_endpoint)
         console.print(f"[green]Embedding verified:[/] dim={dimension}")
@@ -133,6 +145,68 @@ def _download_tier_hint(model_id: str) -> str:
     return "size varies"
 
 
+def _prefetch_with_progress(model_id: str, workspace: Path, *, hf_endpoint: str | None) -> Path:
+    """Download one model while rendering live terminal byte progress.
+
+    Args:
+        model_id: HuggingFace repository ID of the model to download.
+        workspace: Resolved Hebb Mind workspace directory.
+        hf_endpoint: Optional HuggingFace-compatible endpoint.
+
+    Returns:
+        Local directory containing the downloaded model.
+
+    Raises:
+        Exception: Propagates download failures from :func:`prefetch_model`.
+    """
+    size_hint = _download_tier_hint(model_id)
+    columns = (
+        TextColumn("[progress.description]{task.description}"),
+        BarColumn(),
+        TaskProgressColumn(),
+        DownloadColumn(),
+        TransferSpeedColumn(),
+        TimeRemainingColumn(),
+    )
+    with Progress(*columns, console=console, transient=False) as progress:
+        task_id = progress.add_task(
+            f"[cyan]Downloading embedding model[/] [dim]({size_hint})[/]",
+            total=None,
+        )
+
+        def update_progress(bytes_done: int, bytes_total: int, current_file: str) -> None:
+            if _is_file_count_progress(current_file):
+                return
+            description = current_file.strip() or "Downloading embedding model"
+            progress.update(
+                task_id,
+                description=f"[cyan]{description}[/] [dim]({size_hint})[/]",
+                completed=bytes_done,
+                total=bytes_total or None,
+            )
+
+        return prefetch_model(
+            model_id,
+            workspace,
+            hf_endpoint=hf_endpoint,
+            progress_callback=update_progress,
+            suppress_native_progress=True,
+        )
+
+
+def _is_file_count_progress(description: str) -> bool:
+    """Return whether a HuggingFace tqdm event counts files rather than bytes.
+
+    Args:
+        description: Description emitted by the HuggingFace tqdm instance.
+
+    Returns:
+        ``True`` for the snapshot-level ``Fetching N files`` counter.
+    """
+    normalized = description.strip().lower()
+    return normalized.startswith("fetching ") and normalized.endswith(" files")
+
+
 def _persist_region(hf_endpoint: str | None, config_path: Path) -> None:
     update_config_field("hf_endpoint", hf_endpoint or "null", config_path)
 

src/hebb/embedding/catalog.py

@@ -80,7 +80,9 @@ class ProbeResult:
 # Skipping these in ``snapshot_download`` roughly halves the bytes pulled for
 # repos that ship ONNX/OpenVINO/TensorFlow exports alongside the torch weights.
 # NOTE: keep ``*.bin`` and ``*.safetensors`` — ``model_dir_complete`` requires a
-# real weight file (``model.safetensors`` OR ``pytorch_model.bin``).
+# real weight file (``model.safetensors`` OR ``pytorch_model.bin``). The helper
+# below conditionally excludes the exact redundant .bin filename only for
+# built-in repositories verified to publish safetensors.
 PREFETCH_IGNORE_PATTERNS: tuple[str, ...] = (
     "*.onnx",
     "onnx/**",
@@ -93,6 +95,23 @@ class ProbeResult:
     "*.pb",
 )
 
+# These built-in models publish a safetensors checkpoint alongside redundant
+# framework-specific copies. Prefer the safe format while retaining
+# ``pytorch_model.bin`` for BGE-M3 and unknown/custom repositories that may not
+# publish safetensors at all.
+_SAFETENSORS_MODELS = frozenset(
+    {
+        "all-MiniLM-L6-v2",
+        "sentence-transformers/all-MiniLM-L6-v2",
+        "intfloat/multilingual-e5-small",
+        "BAAI/bge-large-en-v1.5",
+    }
+)
+_REDUNDANT_TORCH_WEIGHT_PATTERNS: tuple[str, ...] = (
+    "pytorch_model.bin",
+    "rust_model.ot",
+)
+
 
 def resolve_language(language: str = "auto", environ: dict[str, str] | None = None) -> LanguageSelection:
     """Resolve the content language strategy.
@@ -273,13 +292,14 @@ def prefetch_model(
     workspace: Path,
     hf_endpoint: str | None = None,
     progress_callback: ProgressCallback | None = None,
+    suppress_native_progress: bool = False,
 ) -> Path:
     """Download a HuggingFace model into the Hebb Mind workspace.
 
-    Redundant weight variants (ONNX, OpenVINO, TensorFlow, msgpack) are skipped
-    via :data:`PREFETCH_IGNORE_PATTERNS`; the PyTorch ``*.bin`` / ``*.safetensors``
-    weights plus config, tokenizer, and sentence-transformers module files are
-    always fetched.
+    Redundant weight variants (ONNX, OpenVINO, TensorFlow, msgpack) are skipped.
+    Built-in repositories verified to publish safetensors also skip duplicate
+    ``pytorch_model.bin`` / ``rust_model.ot`` copies, while repositories without
+    safetensors retain their PyTorch ``*.bin`` checkpoint.
 
     Args:
         model_id: HuggingFace repository ID.
@@ -288,6 +308,8 @@ def prefetch_model(
         progress_callback: Optional callback receiving (bytes_done, bytes_total,
             current_file_desc) on every tqdm tick. Used by the web console to
             surface real-time download progress.
+        suppress_native_progress: Hide HuggingFace's tqdm rendering when the
+            callback is displayed by another terminal progress UI.
 
     Returns:
         Local model directory.
@@ -321,12 +343,15 @@ def prefetch_model(
     snapshot_kwargs: dict = {
         "repo_id": model_id,
         "local_dir": str(local_dir),
-        "ignore_patterns": list(PREFETCH_IGNORE_PATTERNS),
+        "ignore_patterns": _prefetch_ignore_patterns(model_id),
     }
     if progress_callback is not None:
         from hebb.embedding.progress import make_progress_tqdm
 
-        snapshot_kwargs["tqdm_class"] = make_progress_tqdm(progress_callback)
+        snapshot_kwargs["tqdm_class"] = make_progress_tqdm(
+            progress_callback,
+            silent=suppress_native_progress,
+        )
 
     try:
         snapshot_download(**snapshot_kwargs)
@@ -342,6 +367,21 @@ def prefetch_model(
     return local_dir
 
 
+def _prefetch_ignore_patterns(model_id: str) -> list[str]:
+    """Return safe redundant-file exclusions for one model repository.
+
+    Args:
+        model_id: HuggingFace repository ID.
+
+    Returns:
+        Ignore patterns that retain at least one supported PyTorch checkpoint.
+    """
+    patterns = list(PREFETCH_IGNORE_PATTERNS)
+    if model_id in _SAFETENSORS_MODELS:
+        patterns.extend(_REDUNDANT_TORCH_WEIGHT_PATTERNS)
+    return patterns
+
+
 def _detect_locale(environ: dict[str, str]) -> str | None:
     for key in ("LC_ALL", "LC_MESSAGES", "LANGUAGE", "LANG"):
         value = environ.get(key)

src/hebb/embedding/progress.py

@@ -20,11 +20,38 @@
 """
 
 
-def make_progress_tqdm(callback: ProgressCallback) -> type:
-    """Return a tqdm subclass that calls ``callback`` on every update."""
+class _SilentProgressOutput:
+    """Discard tqdm rendering while preserving its internal counters."""
+
+    def write(self, text: str) -> int:
+        return len(text)
+
+    def flush(self) -> None:
+        return None
+
+
+_SILENT_PROGRESS_OUTPUT = _SilentProgressOutput()
+
+
+def make_progress_tqdm(callback: ProgressCallback, *, silent: bool = False) -> type:
+    """Return a tqdm subclass that calls ``callback`` on every update.
+
+    Args:
+        callback: Consumer of byte progress updates.
+        silent: Suppress tqdm's own terminal rendering when another UI renders
+            the callback events.
+
+    Returns:
+        A tqdm-compatible progress class.
+    """
     from tqdm import tqdm  # type: ignore[import-untyped]
 
     class ProgressTqdm(tqdm):  # type: ignore[misc]
+        def __init__(self, *args: Any, **kwargs: Any) -> None:
+            if silent:
+                kwargs["file"] = _SILENT_PROGRESS_OUTPUT
+            super().__init__(*args, **kwargs)
+
         def update(self, n: int | float = 1) -> bool | None:
             ret = super().update(n)
             try:

tests/unit/cli/commands/test_setup.py

@@ -23,7 +23,9 @@ def test_setup_initializes_and_selects_english_model(monkeypatch, tmp_path: Path
     monkeypatch.setenv("LANG", "en_US.UTF-8")
     monkeypatch.setattr(
         "hebb.cli.commands.setup.prefetch_model",
-        lambda model_id, workspace, hf_endpoint=None: workspace / "models" / model_id,
+        lambda model_id, workspace, hf_endpoint=None, progress_callback=None, suppress_native_progress=False: (
+            workspace / "models" / model_id
+        ),
     )
     monkeypatch.setattr("hebb.cli.commands.setup._verify_model", lambda model_id, hf_endpoint: 384)
 
@@ -46,7 +48,9 @@ def test_setup_initializes_and_selects_chinese_model(monkeypatch, tmp_path: Path
     monkeypatch.setenv("LANG", "zh_CN.UTF-8")
     monkeypatch.setattr(
         "hebb.cli.commands.setup.prefetch_model",
-        lambda model_id, workspace, hf_endpoint=None: workspace / "models" / model_id,
+        lambda model_id, workspace, hf_endpoint=None, progress_callback=None, suppress_native_progress=False: (
+            workspace / "models" / model_id
+        ),
     )
     monkeypatch.setattr("hebb.cli.commands.setup._verify_model", lambda model_id, hf_endpoint: 384)
 
@@ -68,7 +72,9 @@ def test_setup_best_profile_selects_bge_english(monkeypatch, tmp_path: Path) ->
     monkeypatch.setenv("LANG", "en_US.UTF-8")
     monkeypatch.setattr(
         "hebb.cli.commands.setup.prefetch_model",
-        lambda model_id, workspace, hf_endpoint=None: workspace / "models" / model_id,
+        lambda model_id, workspace, hf_endpoint=None, progress_callback=None, suppress_native_progress=False: (
+            workspace / "models" / model_id
+        ),
     )
     monkeypatch.setattr("hebb.cli.commands.setup._verify_model", lambda model_id, hf_endpoint: 1024)
 
@@ -88,7 +94,9 @@ def test_setup_explicit_language_and_region_are_independent(monkeypatch, tmp_pat
     monkeypatch.setenv("HEBB_HOME", str(home))
     monkeypatch.setattr(
         "hebb.cli.commands.setup.prefetch_model",
-        lambda model_id, workspace, hf_endpoint=None: workspace / "models" / model_id,
+        lambda model_id, workspace, hf_endpoint=None, progress_callback=None, suppress_native_progress=False: (
+            workspace / "models" / model_id
+        ),
     )
     monkeypatch.setattr("hebb.cli.commands.setup._verify_model", lambda model_id, hf_endpoint: 384)
 
@@ -118,7 +126,9 @@ def test_setup_keeps_custom_model_without_explicit_language(monkeypatch, tmp_pat
 
         monkeypatch.setattr(
             "hebb.cli.commands.setup.prefetch_model",
-            lambda model_id, workspace, hf_endpoint=None: workspace / "models" / model_id,
+            lambda model_id, workspace, hf_endpoint=None, progress_callback=None, suppress_native_progress=False: (
+                workspace / "models" / model_id
+            ),
         )
         monkeypatch.setattr("hebb.cli.commands.setup._verify_model", lambda model_id, hf_endpoint: 777)
         result = runner.invoke(setup_cmd, ["--region", "global"])
@@ -155,11 +165,52 @@ def _fail_prefetch(*args: object, **kwargs: object) -> Path:
     assert result.exit_code == 0, result.output
     assert called["prefetch"] is False
     assert "Model already present" in result.output
+    assert "Downloading embedding model" not in result.output
     config = json.loads((home / "hebb.json").read_text())
     assert config["embedding_model"] == "sentence-transformers/all-MiniLM-L6-v2"
     assert config["embedding_dim"] == 384
 
 
+def test_setup_renders_live_download_progress(monkeypatch, tmp_path: Path) -> None:
+    home = tmp_path / "home"
+    monkeypatch.setenv("HEBB_HOME", str(home))
+    _clear_locale_env(monkeypatch)
+    monkeypatch.setenv("LANG", "en_US.UTF-8")
+    monkeypatch.setattr("hebb.cli.commands.setup.workspace_model_available", lambda workspace, model_id: False)
+
+    callback_received = False
+
+    def fake_prefetch(
+        model_id: str,
+        workspace: Path,
+        hf_endpoint: str | None = None,
+        progress_callback=None,
+        suppress_native_progress: bool = False,
+    ) -> Path:
+        nonlocal callback_received
+        assert progress_callback is not None
+        assert suppress_native_progress is True
+        callback_received = True
+        progress_callback(512, 1024, "model.safetensors")
+        progress_callback(1024, 1024, "model.safetensors")
+        progress_callback(12, 12, "Fetching 12 files")
+        return workspace / "models" / model_id
+
+    monkeypatch.setattr("hebb.cli.commands.setup.prefetch_model", fake_prefetch)
+    monkeypatch.setattr("hebb.cli.commands.setup._verify_model", lambda model_id, hf_endpoint: 384)
+
+    runner = CliRunner()
+    with runner.isolated_filesystem(temp_dir=tmp_path):
+        result = runner.invoke(setup_cmd, ["--region", "global"])
+
+    assert result.exit_code == 0, result.output
+    assert callback_received is True
+    assert "model.safetensors" in result.output
+    assert "small ~90MB" in result.output
+    assert "100%" in result.output
+    assert "Fetching 12 files" not in result.output
+
+
 def test_initialize_workspace_uses_hebb_home(monkeypatch, tmp_path: Path) -> None:
     home = tmp_path / "home"
     monkeypatch.setenv("HEBB_HOME", str(home))

tests/unit/embedding/test_catalog.py

@@ -132,12 +132,33 @@ def fake_snapshot_download(**kwargs: object) -> str:
         # Redundant heavy variants are skipped.
         assert "*.onnx" in ignored
         assert "openvino/**" in ignored
-        # The actual weight files must NEVER be excluded — model_dir_complete
-        # requires safetensors OR pytorch_model.bin.
-        assert "*.safetensors" not in ignored
+        # This built-in model ships three equivalent checkpoints; retain only
+        # safetensors so the advertised ~90MB download remains accurate.
+        assert "pytorch_model.bin" in ignored
+        assert "rust_model.ot" in ignored
+        # Never use a broad pattern that would also exclude custom/BGE-M3
+        # repositories whose only supported checkpoint is a .bin file.
         assert "*.bin" not in ignored
+        assert "*.safetensors" not in ignored
         assert "model.safetensors" not in ignored
+
+    def test_prefetch_keeps_bin_for_model_without_safetensors(self, tmp_path: Path, monkeypatch) -> None:
+        captured: dict[str, object] = {}
+
+        def fake_snapshot_download(**kwargs: object) -> str:
+            captured.update(kwargs)
+            return str(kwargs["local_dir"])
+
+        import huggingface_hub
+
+        monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download)
+
+        catalog.prefetch_model("BAAI/bge-m3", tmp_path)
+
+        ignored = captured["ignore_patterns"]
+        assert isinstance(ignored, list)
         assert "pytorch_model.bin" not in ignored
+        assert "*.bin" not in ignored
 
 
 def test_region_auto_prefers_official_when_faster(monkeypatch) -> None: