From 60b37211050a7456d4568ad208f0db7fbb426b0e Mon Sep 17 00:00:00 2001 From: IRENEKO Date: Thu, 5 Nov 2026 19:52:08 +0000 Subject: [PATCH] feat(probe): discover probe slots in adapter libraries A library can now ship a probe as //probe/ with probe_head.safetensors and probe_config.json (labels, layer), so --adapters composes it without a manifest. The slot is named _probe, so it can sit next to a LoRA of the same name. The compose report skips LoRA analysis when every slot is a probe, which used to crash. Signed-off-by: IRENEKO --- .../composer/adapter_discovery.py | 14 ++-- .../composer/compose_granite_switch.py | 12 ++- tests/composer/test_adapter_filtering.py | 16 ++++ tests/composer/test_probe_compose.py | 81 +++++++++++++++++++ tests/composer/test_selective_download.py | 17 ++++ 5 files changed, 134 insertions(+), 6 deletions(-) diff --git a/src/granite_switch/composer/adapter_discovery.py b/src/granite_switch/composer/adapter_discovery.py index 9fd2cd83..4d1cced9 100644 --- a/src/granite_switch/composer/adapter_discovery.py +++ b/src/granite_switch/composer/adapter_discovery.py @@ -60,19 +60,23 @@ def discover_adapters( adapter_dir = io_yaml_path.parent tech = adapter_dir.name - if tech not in ("alora", "lora"): + if tech not in ("alora", "lora", "probe"): tech = technology_fallback if not tech or (technology_filter and tech != technology_filter): continue model_name = adapter_dir.parent.name adapter_name = adapter_dir.parent.parent.name + if tech == "probe": + adapter_name += "_probe" if model_name == target_model_name: - adapter_model = adapter_dir / "adapter_model.safetensors" - adapter_config_file = adapter_dir / "adapter_config.json" + if tech == "probe": + artifacts = ("probe_head.safetensors", "probe_config.json") + else: + artifacts = ("adapter_model.safetensors", "adapter_config.json") - if adapter_model.exists() and adapter_config_file.exists(): + if all((adapter_dir / f).exists() for f in artifacts): if adapter_name in discovered_by_name: existing_tech = discovered_by_name[adapter_name][2] if tech == "alora" and existing_tech == "lora": @@ -101,7 +105,7 @@ def discover_adapters( print(f"\nDiscovery complete: {len(discovered)} adapter(s) found") # Report per-adapter module contributions - _report_module_contributions(discovered, arch) + _report_module_contributions([d for d in discovered if d[2] != "probe"], arch) return discovered diff --git a/src/granite_switch/composer/compose_granite_switch.py b/src/granite_switch/composer/compose_granite_switch.py index d7971982..1ed315e1 100755 --- a/src/granite_switch/composer/compose_granite_switch.py +++ b/src/granite_switch/composer/compose_granite_switch.py @@ -1035,6 +1035,15 @@ def build(): include=args.include_adapters, exclude=args.exclude_adapters, ) + for path, name, tech, _source in found: + if tech == "probe": + with open(Path(path) / "probe_config.json") as f: + cfg = json.load(f) + probe_meta[name] = { + "kind": "probe", + "labels": cfg.get("labels"), + "layer": cfg.get("layer"), + } if not found: msg = ( f" WARNING: No adapters found for target '{args.target_model}'" @@ -1379,12 +1388,13 @@ def build(): print("Generating compose report...") print("=" * 80) if hasattr(model, "_build_mappings"): + has_lora = any(t[0] and t[2] != ANCHOR_MODE_PROBE for t in all_discovered) generate_compose_report( base_mapping=model._build_mappings["base"], adapter_mapping=model._build_mappings["adapter"], output_path=args.output, model=model, - adapter_paths=adapter_paths, + adapter_paths=adapter_paths if has_lora else None, adapter_names=adapter_names, arch=arch, source_analysis=model._build_mappings.get("source_analysis"), diff --git a/tests/composer/test_adapter_filtering.py b/tests/composer/test_adapter_filtering.py index d1b641c7..491233c9 100644 --- a/tests/composer/test_adapter_filtering.py +++ b/tests/composer/test_adapter_filtering.py @@ -132,6 +132,22 @@ def test_filter_and_override(self, tmp_path): # -- filter_adapters tests --------------------------------------------------- +class TestProbeDiscovery: + def test_probe_dir_is_discovered_with_suffix(self, tmp_path): + root = _make_adapter_library(tmp_path, "granite-4.0-micro", [("ans", ["lora"])]) + probe = tmp_path / "ans" / "granite-4.0-micro" / "probe" + probe.mkdir(parents=True) + (probe / "io.yaml").write_text("name: ans_probe\n") + (probe / "probe_config.json").write_text('{"labels": ["no", "yes"]}') + (probe / "probe_head.safetensors").write_bytes(b"\x00") + arch = resolve_arch("ibm-granite/granite-4.0-micro") + found = discover_adapters(root, "granite-4.0-micro", arch) + assert sorted((n, t) for _, n, t, _ in found) == [ + ("ans", "lora"), + ("ans_probe", "probe"), + ] + + class TestFilterAdapters: def test_no_filters_passthrough(self, sample_discovered): result = filter_adapters(sample_discovered) diff --git a/tests/composer/test_probe_compose.py b/tests/composer/test_probe_compose.py index 0af60d29..e2aa6a91 100644 --- a/tests/composer/test_probe_compose.py +++ b/tests/composer/test_probe_compose.py @@ -513,3 +513,84 @@ def test_probe_substitute_is_the_bases_role_open_marker(base, role_open): ) assert substitutes == [tokenizer.convert_tokens_to_ids(role_open)] + + +def _compose_cli(base, library, output): + import subprocess + import sys + + cmd = [ + sys.executable, + "-m", + "granite_switch.composer.compose_granite_switch", + "--base-model", + base, + "--adapters", + str(library), + "--output", + str(output), + ] + result = subprocess.run(cmd, capture_output=True, text=True, timeout=1800) + assert result.returncode == 0, result.stdout[-3000:] + result.stderr[-3000:] + with open(output / "config.json") as f: + return json.load(f) + + +def _write_library_probe(library, name, target, labels, layer, hidden): + probe_dir = library / name / target / "probe" + _write_synthetic_head( + str(probe_dir), *_winning_head(len(labels), hidden, 0), labels + ) + (probe_dir / "probe_config.json").write_text( + json.dumps({"labels": labels, "layer": layer}) + ) + (probe_dir / "io.yaml").write_text(f"name: {name}\nmodel: ~\n") + + +def test_probe_only_library_compose_cli(tmp_path): + """`--adapters ` composes `//probe/` slots alone.""" + tokenizer = _tokenizer() + labels, label_ids = _single_token_labels( + tokenizer, ["safe", "unsafe", "yes", "no", "true", "false"] + ) + library = tmp_path / "lib" + _write_library_probe( + library, "detect", "granite-4.0-micro", labels, 12, _hidden_size() + ) + + cfg = _compose_cli(BASE_MODEL, library, tmp_path / "out") + + assert cfg["adapter_names"] == ["detect_probe"] + assert cfg["adapter_kinds"] == ["probe"] + assert cfg["probe_label_token_ids"] == [label_ids] + assert cfg["probe_read_layers"] == [12] + assert (tmp_path / "out" / "io_configs" / "detect_probe" / "io.yaml").is_file() + + +def test_library_with_lora_and_probe_compose_cli(tmp_path): + """A LoRA and a probe under one name compose as two slots, LoRA first.""" + import shutil + + tokenizer = _tokenizer(MIXED_BASE) + labels, label_ids = _single_token_labels( + tokenizer, ["safe", "unsafe", "yes", "no", "true", "false"] + ) + library = tmp_path / "lib" + shutil.copytree( + _find_rag_lora_dir(), library / "answerability" / MIXED_LORA_TARGET / "lora" + ) + _write_library_probe( + library, + "answerability", + MIXED_LORA_TARGET, + labels, + None, + _hidden_size(MIXED_BASE), + ) + + cfg = _compose_cli(MIXED_BASE, library, tmp_path / "out") + + assert cfg["adapter_names"] == ["answerability", "answerability_probe"] + assert cfg["adapter_kinds"] == ["lora", "probe"] + assert cfg["probe_label_token_ids"] == [None, label_ids] + assert cfg["probe_read_layers"] == [None, None] diff --git a/tests/composer/test_selective_download.py b/tests/composer/test_selective_download.py index 64a95f7f..664d7cb8 100644 --- a/tests/composer/test_selective_download.py +++ b/tests/composer/test_selective_download.py @@ -190,6 +190,23 @@ def test_technology_filter_overrides_alora_preference(self, default_tree): "query_rewrite/granite-4.1-3b/lora/**", ] + def test_probe_dir_is_downloaded(self): + tree = _tree_response( + { + "": ["answerability", "safety"], + "answerability/granite-4.1-3b": ["lora"], + "safety/granite-4.1-3b": ["probe"], + } + ) + with patch("huggingface_hub.list_repo_tree", side_effect=tree): + patterns = _build_allow_patterns( + "org/repo", target_model_name="granite-4.1-3b" + ) + assert patterns == [ + "answerability/granite-4.1-3b/lora/**", + "safety/granite-4.1-3b/**", + ] + # --------------------------------------------------------------------------- # resolve_repo_path — integration with snapshot_download