diff --git a/src/semble/index/dense.py b/src/semble/index/dense.py index f7f91e2a..ada1f875 100644 --- a/src/semble/index/dense.py +++ b/src/semble/index/dense.py @@ -1,5 +1,6 @@ from __future__ import annotations +import logging from functools import cache from pathlib import Path @@ -15,11 +16,17 @@ from semble.utils import resolve_model_name +def _drop_unauthenticated_warning(record: logging.LogRecord) -> bool: + """Drop the Hub's unauthenticated-request warning; the public model downloads fine without a token.""" + return "unauthenticated requests" not in record.getMessage() + + @cache def _load_cached(model_path: str) -> StaticModel: """Load a model and cache it, but only after the path resolves.""" # Disable HF progress bars since the model is loaded silently in the background during indexing. disable_progress_bars() + logging.getLogger("huggingface_hub.utils._http").addFilter(_drop_unauthenticated_warning) try: try: model = StaticModel.from_pretrained(model_path, force_download=False) diff --git a/tests/test_search.py b/tests/test_search.py index fa3d3bd1..d17a7d8e 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -1,3 +1,4 @@ +import logging from typing import Any from unittest.mock import MagicMock, call, patch @@ -159,6 +160,28 @@ def test_load_model(model_path: str | None, expected_call_arg: str, incomplete_c assert mock_fp.call_args_list == expected_calls +@pytest.mark.parametrize( + ("message", "shown"), + [ + ( + "Warning: You are sending unauthenticated requests to the HF Hub. " + "Please set a HF_TOKEN to enable higher rate limits and faster downloads.", + False, + ), + ("Rate limited. Waiting 1s before retry [Retry 1/5].", True), + ("Your HF_TOKEN expires soon.", True), + ], +) +def test_load_model_hides_only_unauthenticated_warning( + caplog: pytest.LogCaptureFixture, message: str, shown: bool +) -> None: + """Loading the model hides the Hub's unauthenticated-request nag but keeps other Hub warnings.""" + with patch("semble.index.dense.StaticModel.from_pretrained"): + load_model(f"filter/test-{shown}") + logging.getLogger("huggingface_hub.utils._http").warning(message) + assert (message in caplog.text) is shown + + def test_embed_chunks_empty_returns_empty_array(mock_model: Any) -> None: """embed_chunks with an empty list returns a (0, 256) float32 array.""" result = embed_chunks(mock_model, [])