feat(probe): discover probe slots in adapter libraries

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
 

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"),

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)

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 <library>` composes `<name>/<model>/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]

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