diff --git a/src/granite_switch/composer/adapter_discovery.py b/src/granite_switch/composer/adapter_discovery.py index 7c175ca..75ad1f8 100644 --- a/src/granite_switch/composer/adapter_discovery.py +++ b/src/granite_switch/composer/adapter_discovery.py @@ -279,12 +279,182 @@ def is_adapter_library(path: str) -> bool: return p.is_dir() and not (p / "adapter_config.json").exists() -def resolve_repo_path(path_or_repo: str) -> str: +# ------------------------------------------------------------------ # +# HuggingFace Hub metadata helpers (no file downloads) +# ------------------------------------------------------------------ # + + +def _list_repo_adapter_names(repo_id: str) -> List[str]: + """Get adapter folder names from a HF repo using metadata-only API calls. + + Returns top-level directory names, skipping entries that start with ``_`` + (e.g. ``_ollama``). + """ + from huggingface_hub import list_repo_tree + from huggingface_hub.hf_api import RepoFolder + + tree = list_repo_tree(repo_id, repo_type="model") + return [ + item.path for item in tree + if isinstance(item, RepoFolder) and not item.path.startswith("_") + ] + + +def _resolve_technology( + repo_id: str, + adapter_name: str, + target_model_name: str, +) -> Optional[str]: + """Resolve preferred technology for an adapter via Hub metadata. + + Prefers ``alora`` over ``lora``. Returns ``None`` if neither exists + for this adapter/model combination. + """ + from huggingface_hub import list_repo_tree + from huggingface_hub.hf_api import RepoFolder + from huggingface_hub.errors import EntryNotFoundError + + try: + subtree = list_repo_tree( + repo_id, repo_type="model", + path_in_repo=f"{adapter_name}/{target_model_name}", + ) + technologies = { + item.path.split("/")[-1] for item in subtree + if isinstance(item, RepoFolder) + and item.path.split("/")[-1] in ("alora", "lora") + } + except EntryNotFoundError: + return None + + if "alora" in technologies: + return "alora" + elif "lora" in technologies: + return "lora" + return None + + +def _build_allow_patterns( + repo_id: str, + target_model_name: Optional[str] = None, + include_adapters: Optional[List[str]] = None, + exclude_adapters: Optional[List[str]] = None, +) -> Optional[List[str]]: + """Build ``allow_patterns`` for selective ``snapshot_download``. + + Uses lightweight Hub API calls to discover adapter names, then applies + fnmatch-based include/exclude filtering and target model constraints to + construct download patterns. When both adapter names and target model + are known, also resolves the preferred technology (alora > lora) so + only the needed technology variant is downloaded. + + Returns: + List of glob patterns, or ``None`` if no filtering is possible. + """ + adapter_names = _list_repo_adapter_names(repo_id) + + # Apply include filter + if include_adapters: + adapter_names = [ + name for name in adapter_names + if any(fnmatch(name, pat) for pat in include_adapters) + ] + + # Apply exclude filter + if exclude_adapters: + adapter_names = [ + name for name in adapter_names + if not any(fnmatch(name, pat) for pat in exclude_adapters) + ] + + # Construct patterns with technology resolution + if adapter_names and target_model_name: + patterns = [] + for name in adapter_names: + tech = _resolve_technology(repo_id, name, target_model_name) + if tech: + patterns.append(f"{name}/{target_model_name}/{tech}/**") + else: + # Model not found for this adapter — include anyway so + # discover_adapters can report it as missing downstream + patterns.append(f"{name}/{target_model_name}/**") + return patterns if patterns else None + elif target_model_name: + return [f"*/{target_model_name}/**"] + elif adapter_names: + return [f"{name}/**" for name in adapter_names] + else: + return None + + +def list_repo_adapters_remote( + repo_id: str, + target_model_name: str, +) -> List[Dict[str, object]]: + """List adapters available in a remote HF repo without downloading. + + Uses Hub metadata API calls to discover adapter names and their + available technologies for the given target model. + + Args: + repo_id: HuggingFace repo ID (e.g., ``"ibm-granite/granitelib-rag-r1.0"``). + target_model_name: Target model name (e.g., ``"granite-4.1-3b"``). + + Returns: + List of dicts ``{"name": str, "technologies": [str]}``, sorted + by adapter name. + """ + from huggingface_hub import list_repo_tree + from huggingface_hub.hf_api import RepoFolder + from huggingface_hub.errors import EntryNotFoundError + + adapter_names = _list_repo_adapter_names(repo_id) + results = [] + + for name in adapter_names: + try: + subtree = list_repo_tree( + repo_id, repo_type="model", + path_in_repo=f"{name}/{target_model_name}", + ) + technologies = sorted( + item.path.split("/")[-1] for item in subtree + if isinstance(item, RepoFolder) + and item.path.split("/")[-1] in ("alora", "lora") + ) + if technologies: + results.append({"name": name, "technologies": technologies}) + except EntryNotFoundError: + # Adapter doesn't have this target model — skip + continue + + return sorted(results, key=lambda x: x["name"]) + + +# ------------------------------------------------------------------ # +# Path resolution +# ------------------------------------------------------------------ # + + +def resolve_repo_path( + path_or_repo: str, + target_model_name: Optional[str] = None, + include_adapters: Optional[List[str]] = None, + exclude_adapters: Optional[List[str]] = None, +) -> str: """Resolve a local path or HuggingFace repo ID to a local directory. + For HuggingFace repos, applies selective downloading using + ``allow_patterns`` constructed from the provided filters. + Args: path_or_repo: Either a local directory path or a HuggingFace repo ID (e.g., ``"ibm-granite/granite-lib-rag-r1.0"``). + target_model_name: Target model name to filter by (e.g., + ``"granite-4.1-3b"``). + include_adapters: Only download adapters matching these fnmatch + patterns. + exclude_adapters: Skip adapters matching these fnmatch patterns. Returns: Absolute local path to the directory. @@ -298,12 +468,30 @@ def resolve_repo_path(path_or_repo: str) -> str: if "/" in path_or_repo and not local.exists(): print(f" Detected HuggingFace repo: {path_or_repo}") + + # Build selective download patterns + allow_patterns = None + if target_model_name or include_adapters or exclude_adapters: + try: + allow_patterns = _build_allow_patterns( + path_or_repo, + target_model_name=target_model_name, + include_adapters=include_adapters, + exclude_adapters=exclude_adapters, + ) + if allow_patterns: + print(f" Selective download patterns: {allow_patterns}") + except Exception as e: + print(f" WARNING: Failed to build download filters ({e}), " + f"downloading full repo") + allow_patterns = None + print(f" Downloading from HuggingFace Hub...") try: - cache_dir = snapshot_download( - repo_id=path_or_repo, - repo_type="model", - ) + kwargs = {"repo_id": path_or_repo, "repo_type": "model"} + if allow_patterns: + kwargs["allow_patterns"] = allow_patterns + cache_dir = snapshot_download(**kwargs) print(f" Downloaded to: {cache_dir}") return cache_dir except Exception as e: diff --git a/src/granite_switch/composer/compose_granite_switch.py b/src/granite_switch/composer/compose_granite_switch.py index f67068a..4d321a1 100755 --- a/src/granite_switch/composer/compose_granite_switch.py +++ b/src/granite_switch/composer/compose_granite_switch.py @@ -56,6 +56,7 @@ filter_adapters, is_adapter_library, list_available_adapters, + list_repo_adapters_remote, resolve_repo_path, ) from granite_switch.composer.tokenizer_setup import ( @@ -530,15 +531,30 @@ def build(): print("ERROR: --list-adapters requires --adapters") return 1 for entry in args.adapters: - try: - resolved_path = resolve_repo_path(entry) - except Exception as e: - print(f"Failed to resolve {entry}: {e}") - return 1 - if not is_adapter_library(resolved_path): - print(f"\n{entry} is a single adapter, not a library.") - continue - available = list_available_adapters(resolved_path, args.target_model) + # For HF repos, use metadata-only listing (no download) + local = Path(entry) + if "/" in entry and not local.exists(): + try: + available = list_repo_adapters_remote( + entry, args.target_model + ) + except Exception as e: + print(f"Failed to list adapters from {entry}: {e}") + return 1 + else: + # Local path — resolve and scan + try: + resolved_path = resolve_repo_path(entry) + except Exception as e: + print(f"Failed to resolve {entry}: {e}") + return 1 + if not is_adapter_library(resolved_path): + print(f"\n{entry} is a single adapter, not a library.") + continue + available = list_available_adapters( + resolved_path, args.target_model + ) + if not available: print(f"\nNo adapters found in {entry} for target '{args.target_model}'") continue @@ -587,7 +603,12 @@ def build(): for entry in args.adapters: print(f"\nResolving: {entry}") try: - resolved_path = resolve_repo_path(entry) + resolved_path = resolve_repo_path( + entry, + target_model_name=args.target_model, + include_adapters=args.include_adapters, + exclude_adapters=args.exclude_adapters, + ) except Exception as e: print(f"Failed to resolve {entry}: {e}") return 1 diff --git a/tests/composer/test_selective_download.py b/tests/composer/test_selective_download.py new file mode 100644 index 0000000..7386e9e --- /dev/null +++ b/tests/composer/test_selective_download.py @@ -0,0 +1,300 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for selective HuggingFace download behavior. + +Covers the metadata-first path introduced to fix #3: adapter libraries are +filtered *before* download via ``allow_patterns`` derived from +``--base-model``, ``--include-adapters``, and ``--exclude-adapters``. HF Hub +calls (``list_repo_tree``, ``snapshot_download``) are mocked so tests run +offline, except the ``TestRealHubMetadata`` class which hits the real Hub. +""" + +from unittest.mock import patch, MagicMock + +import pytest +from huggingface_hub.hf_api import RepoFile, RepoFolder +from huggingface_hub.errors import EntryNotFoundError + +from granite_switch.composer.adapter_discovery import ( + _build_allow_patterns, + _resolve_technology, + list_repo_adapters_remote, + resolve_repo_path, +) + + +# --------------------------------------------------------------------------- +# Mock helpers +# --------------------------------------------------------------------------- + + +def _folder(path): + return RepoFolder(path=path, oid="deadbeef") + + +def _file(path): + return RepoFile(path=path, oid="deadbeef", size=1) + + +def _tree_response(tree_map): + """Return a ``list_repo_tree`` side_effect that serves *tree_map*. + + *tree_map* maps ``path_in_repo`` (``""`` for root) to the list of + folder/file names that live at that path. Unknown paths raise + ``EntryNotFoundError`` — matching real HF behavior for missing folders. + """ + def _side_effect(repo_id, repo_type="model", path_in_repo=None): + key = path_in_repo or "" + if key not in tree_map: + raise EntryNotFoundError(f"Path not found: {key}") + entries = [] + for entry in tree_map[key]: + full_path = f"{key}/{entry}" if key else entry + if entry.endswith(".txt") or entry.endswith(".md") or entry.endswith(".json"): + entries.append(_file(full_path)) + else: + entries.append(_folder(full_path)) + return entries + return _side_effect + + +# --------------------------------------------------------------------------- +# _resolve_technology — alora is preferred over lora +# --------------------------------------------------------------------------- + + +class TestResolveTechnology: + def test_prefers_alora_when_both_exist(self): + tree = _tree_response({ + "answerability/granite-4.1-3b": ["alora", "lora"], + }) + with patch("huggingface_hub.list_repo_tree", side_effect=tree): + assert _resolve_technology( + "org/repo", "answerability", "granite-4.1-3b" + ) == "alora" + + def test_returns_none_when_target_model_missing(self): + tree = _tree_response({}) + with patch("huggingface_hub.list_repo_tree", side_effect=tree): + assert _resolve_technology( + "org/repo", "answerability", "granite-99b" + ) is None + + +# --------------------------------------------------------------------------- +# _build_allow_patterns — pattern construction from filters +# --------------------------------------------------------------------------- + + +class TestBuildAllowPatterns: + def _default_tree(self): + return _tree_response({ + "": ["answerability", "citations", "query_rewrite"], + "answerability/granite-4.1-3b": ["alora", "lora"], + "citations/granite-4.1-3b": ["lora"], + "query_rewrite/granite-4.1-3b": ["alora"], + }) + + def test_all_adapters_with_target_model(self): + with patch( + "huggingface_hub.list_repo_tree", + side_effect=self._default_tree(), + ): + patterns = _build_allow_patterns( + "org/repo", target_model_name="granite-4.1-3b", + ) + assert patterns == [ + "answerability/granite-4.1-3b/alora/**", + "citations/granite-4.1-3b/lora/**", + "query_rewrite/granite-4.1-3b/alora/**", + ] + + def test_include_filter_applies(self): + with patch( + "huggingface_hub.list_repo_tree", + side_effect=self._default_tree(), + ): + patterns = _build_allow_patterns( + "org/repo", + target_model_name="granite-4.1-3b", + include_adapters=["answerability"], + ) + assert patterns == ["answerability/granite-4.1-3b/alora/**"] + + +# --------------------------------------------------------------------------- +# resolve_repo_path — integration with snapshot_download +# --------------------------------------------------------------------------- + + +class TestResolveRepoPathSelectiveDownload: + def test_local_path_returns_as_is_without_download(self, tmp_path): + (tmp_path / "adapter_config.json").write_text("{}") + with patch("huggingface_hub.snapshot_download") as mock_dl: + result = resolve_repo_path( + str(tmp_path), + target_model_name="granite-4.1-3b", + include_adapters=["anything"], + ) + mock_dl.assert_not_called() + assert result == str(tmp_path) + + def test_hf_repo_passes_allow_patterns(self, tmp_path): + tree = _tree_response({ + "": ["answerability", "citations"], + "answerability/granite-4.1-3b": ["alora"], + "citations/granite-4.1-3b": ["lora"], + }) + mock_dl = MagicMock(return_value=str(tmp_path)) + with patch( + "huggingface_hub.list_repo_tree", side_effect=tree, + ), patch("huggingface_hub.snapshot_download", mock_dl): + resolve_repo_path( + "org/repo", target_model_name="granite-4.1-3b", + ) + kwargs = mock_dl.call_args.kwargs + assert kwargs["allow_patterns"] == [ + "answerability/granite-4.1-3b/alora/**", + "citations/granite-4.1-3b/lora/**", + ] + + def test_hf_repo_without_filters_downloads_full(self, tmp_path): + """No filters → allow_patterns not passed (matches pre-fix behavior).""" + mock_dl = MagicMock(return_value=str(tmp_path)) + with patch("huggingface_hub.snapshot_download", mock_dl): + resolve_repo_path("org/repo") + assert "allow_patterns" not in mock_dl.call_args.kwargs + + def test_pattern_build_failure_falls_back_to_full_download(self, tmp_path): + """If metadata pass raises, warn and continue with a full download.""" + def _boom(*args, **kwargs): + raise RuntimeError("HF Hub down") + + mock_dl = MagicMock(return_value=str(tmp_path)) + with patch( + "huggingface_hub.list_repo_tree", side_effect=_boom, + ), patch("huggingface_hub.snapshot_download", mock_dl): + resolve_repo_path( + "org/repo", target_model_name="granite-4.1-3b", + ) + assert "allow_patterns" not in mock_dl.call_args.kwargs + + +# --------------------------------------------------------------------------- +# End-to-end: exact scenario from issue #3 +# --------------------------------------------------------------------------- + + +class TestIssue3ReproScenario: + """Mirrors the "Steps to reproduce" from issue #3 verbatim. + + Invocation:: + + python -m granite_switch.composer.compose_granite_switch \\ + --base-model ibm-granite/granite-4.1-3b \\ + --adapters ibm-granite/granitelib-core-r1.0 \\ + ibm-granite/granitelib-rag-r1.0 \\ + --include-adapters query_rewrite context-attribution + + Verifies that each ``snapshot_download`` call carries ``allow_patterns`` + restricted to exactly the requested adapter/model paths — no 8b/30b + variants, only the correct technology per adapter. + """ + + def _core_tree(self): + return _tree_response({ + "": ["context-attribution", "requirement-check", "uncertainty"], + "context-attribution/granite-4.1-3b": ["lora"], + "context-attribution/granite-4.1-8b": ["lora"], + "requirement-check/granite-4.1-3b": ["alora"], + "uncertainty/granite-4.1-3b": ["alora"], + }) + + def _rag_tree(self): + return _tree_response({ + "": [ + "query_rewrite", "answerability", + "citations", "hallucination_detection", + ], + "query_rewrite/granite-4.1-3b": ["alora"], + "query_rewrite/granite-4.1-8b": ["alora"], + "answerability/granite-4.1-3b": ["alora"], + "citations/granite-4.1-3b": ["lora"], + "hallucination_detection/granite-4.1-3b": ["lora"], + }) + + def test_only_requested_adapter_model_paths_are_downloaded(self, tmp_path): + target_model = "granite-4.1-3b" + include = ["query_rewrite", "context-attribution"] + + mock_dl = MagicMock(return_value=str(tmp_path)) + + with patch( + "huggingface_hub.list_repo_tree", side_effect=self._core_tree(), + ), patch("huggingface_hub.snapshot_download", mock_dl): + resolve_repo_path( + "ibm-granite/granitelib-core-r1.0", + target_model_name=target_model, + include_adapters=include, + ) + assert mock_dl.call_args.kwargs["allow_patterns"] == [ + "context-attribution/granite-4.1-3b/lora/**", + ] + + mock_dl.reset_mock() + with patch( + "huggingface_hub.list_repo_tree", side_effect=self._rag_tree(), + ), patch("huggingface_hub.snapshot_download", mock_dl): + resolve_repo_path( + "ibm-granite/granitelib-rag-r1.0", + target_model_name=target_model, + include_adapters=include, + ) + assert mock_dl.call_args.kwargs["allow_patterns"] == [ + "query_rewrite/granite-4.1-3b/alora/**", + ] + + +# --------------------------------------------------------------------------- +# Real HuggingFace Hub metadata calls (no file downloads) +# --------------------------------------------------------------------------- +# +# Uses the real HF Hub API via ``list_repo_tree`` to verify the helpers +# against an actual repo layout. ``snapshot_download`` is still mocked so +# nothing heavy is pulled to disk. Marked ``slow`` because they require +# network; skip with ``-m "not slow"``. + + +@pytest.mark.slow +class TestRealHubMetadata: + REPO = "ibm-granite/granitelib-core-r1.0" + TARGET_MODEL = "granite-4.1-3b" + + def test_resolve_technology_matches_published_build(self): + # context-attribution is documented as 'lora' in the published BUILD.md; + # requirement-check as 'alora'. + assert _resolve_technology( + self.REPO, "context-attribution", self.TARGET_MODEL, + ) == "lora" + assert _resolve_technology( + self.REPO, "requirement-check", self.TARGET_MODEL, + ) == "alora" + + def test_list_repo_adapters_remote_includes_known_adapters(self): + known = {"context-attribution", "requirement-check", "uncertainty"} + result = list_repo_adapters_remote(self.REPO, self.TARGET_MODEL) + names = {entry["name"] for entry in result} + assert known.issubset(names), ( + f"Missing adapters. Expected ⊇ {known}, got {names}" + ) + + def test_build_allow_patterns_against_real_repo(self, tmp_path): + mock_dl = MagicMock(return_value=str(tmp_path)) + with patch("huggingface_hub.snapshot_download", mock_dl): + resolve_repo_path( + self.REPO, + target_model_name=self.TARGET_MODEL, + include_adapters=["context-attribution"], + ) + assert mock_dl.call_args.kwargs["allow_patterns"] == [ + "context-attribution/granite-4.1-3b/lora/**", + ]