feat(classifier): add switched classifier slots

src/granite_switch/composer/adapter_discovery.py

@@ -108,24 +108,48 @@ def discover_adapters(
 
 def discover_adapters_from_yaml(
     manifest_path: str,
-) -> list[tuple[str, str, str, str | None]]:
+) -> tuple[list[tuple[str, str, str, str | None]], dict[str, dict]]:
     """Discover adapters from a YAML manifest file.
 
     Reads a YAML manifest that maps adapter names to their paths and types.
+    An entry can also mark itself a **probe** slot with these extra keys:
+
+    .. code-block:: yaml
+
+        hallucination:
+          path: /path/to/head
+          type: lora                # lora/alora technology (token placement)
+          kind: probe               # default "lora" when omitted
+          labels: [faithful, hallucinated]
+          layer: 12                 # optional decoder layer; omit for the final layer
+
+    ``type`` and ``kind`` are independent axes: ``type`` is the LoRA technology
+    (lora vs alora, which drives token placement), ``kind`` is whether the slot
+    is a real LoRA or a probe head. ``kind`` and ``labels`` are optional;
+    without them an entry is a plain LoRA adapter. The label words become token
+    ids at compose time (see :func:`tokenizer_setup.resolve_label_token_ids`).
+
+    ``layer: K`` reads the residual stream after zero-based decoder layer ``K``
+    (switch cache layers excluded). Omitting it reads the final post-norm state.
 
     Args:
         manifest_path: Path to the YAML manifest file.
 
     Returns:
+        ``(found, probe_meta)``. ``found`` is:
         List of ``(adapter_path, adapter_name, technology, source)`` tuples.
         The source is set to the manifest path for traceability.
+        ``probe_meta`` maps name ->
+        ``{"kind", "labels", "layer"}`` for each entry that set any of them;
+        names not in it are plain LoRA slots.
     """
     import yaml
 
     path = Path(manifest_path)
     print(f"  Loading adapters manifest: {path.name}")
 
     found = []
+    probe_meta: dict[str, dict] = {}
     if path.is_file() and path.suffix in (".yaml", ".yml"):
         with open(path) as f:
             adapters_config = yaml.safe_load(f)
@@ -137,7 +161,18 @@ def discover_adapters_from_yaml(
                 # Use manifest path as source for traceability
                 found.append((adapter_path, name, tech, manifest_path))
 
-    return found
+                # Optional probe slots.
+                kind = info.get("kind")
+                labels = info.get("labels")
+                layer = info.get("layer")
+                if kind is not None or labels is not None or layer is not None:
+                    probe_meta[name] = {
+                        "kind": kind or "lora",
+                        "labels": labels,
+                        "layer": layer,
+                    }
+
+    return found, probe_meta
 
 
 def _report_module_contributions(

src/granite_switch/composer/arch.py

@@ -112,12 +112,14 @@ class ArchDescriptor:
     # LoRA parameter keywords (for filtering)
     lora_keywords: list[str] = field(default_factory=lambda: ["lora_A", "lora_B"])
 
-    # Non-LoRA buffer keywords to exclude from base validation
+    # Non-LoRA/probe buffer keywords to exclude from base validation
     buffer_keywords: list[str] = field(
         default_factory=lambda: [
             "adapter_token_ids",
             "adapter_scalings",
             "control_to_substitute_lut",
+            "probe_head",
+            "adapter_kind_lut",
         ]
     )
 
@@ -392,14 +394,6 @@ def granite_dense_arch(base_config=None) -> ArchDescriptor:
     )
 
 
-_SR_BUFFER_KEYWORDS: list[str] = [
-    "adapter_token_ids",
-    "adapter_scalings",
-    "control_to_substitute_lut",
-    "cross_stream",  # zero-initialized, no base source
-]
-
-
 def granite_moe_hybrid_sr_arch(base_config=None) -> ArchDescriptor:
     """GraniteMoeHybrid Shadow Residual architecture.
 
@@ -408,23 +402,23 @@ def granite_moe_hybrid_sr_arch(base_config=None) -> ArchDescriptor:
     """
     arch = granite_moe_hybrid_arch(base_config=base_config)
     arch.groups = arch.groups + _cross_stream_groups()
-    arch.buffer_keywords = list(_SR_BUFFER_KEYWORDS)
+    arch.buffer_keywords.append("cross_stream")  # zero-initialized, no base source
     return arch
 
 
 def granite_moe_sr_arch(base_config=None) -> ArchDescriptor:
     """GraniteMoe Shadow Residual architecture (attention + ``cross_stream``)."""
     arch = granite_moe_arch(base_config=base_config)
     arch.groups = arch.groups + _cross_stream_groups()
-    arch.buffer_keywords = list(_SR_BUFFER_KEYWORDS)
+    arch.buffer_keywords.append("cross_stream")
     return arch
 
 
 def granite_dense_sr_arch(base_config=None) -> ArchDescriptor:
     """Granite dense Shadow Residual architecture (fused + ``cross_stream``)."""
     arch = granite_dense_arch(base_config=base_config)
     arch.groups = arch.groups + _cross_stream_groups()
-    arch.buffer_keywords = list(_SR_BUFFER_KEYWORDS)
+    arch.buffer_keywords.append("cross_stream")
     return arch
 
 

src/granite_switch/composer/compose_granite_switch.py

@@ -69,14 +69,17 @@
 from granite_switch.composer.compose_utils import GraniteSwitchComposer
 from granite_switch.composer.reporting import generate_compose_report, write_build_doc
 from granite_switch.composer.tokenizer_setup import (
+    ANCHOR_MODE_PROBE,
     ANCHOR_MODE_SR,
     add_audio_token,
     add_control_tokens,
     build_substitute_token_ids,
     configure_audio_chat_template,
     configure_chat_template,
+    detect_template_format,
     find_reserved_never_emitted_token_id,
     load_activation_anchor,
+    resolve_label_token_ids,
 )
 from granite_switch.composer.validator import validate_control_lut
 from granite_switch.composer.weight_transfer import validate_untied_lm_head_saved
@@ -119,10 +122,26 @@ def build_control_token_lists(tokenizer, all_discovered, base_reset):
         all_discovered,
         _probe_lora_substitute_token_id(tokenizer),
         base_reset=base_reset,
+        probe_slot_substitute_id=_probe_slot_substitute_token_id(tokenizer),
     )
     return adapter_token_ids, special_tokens, adapter_substitute_token_ids
 
 
+def _probe_slot_substitute_token_id(tokenizer) -> int:
+    """The base's role-open marker id, which follows the last turn's close."""
+    fmt = detect_template_format(tokenizer.chat_template)
+    if fmt is None:
+        raise ValueError(
+            "Unrecognized chat template; cannot pick the probe substitute."
+        )
+    ids = tokenizer.encode(fmt.role_open_marker, add_special_tokens=False)
+    if len(ids) != 1:
+        raise ValueError(
+            f"Role-open marker {fmt.role_open_marker!r} encodes to {len(ids)} tokens."
+        )
+    return ids[0]
+
+
 def _probe_lora_substitute_token_id(tokenizer) -> int:
     """Ask the tokenizer which token naturally appears at sequence position 0
     of a rendered no-adapter chat.
@@ -981,6 +1000,7 @@ def build():
     step_start = time.time()
 
     discovered_adapters = []
+    probe_meta: dict[str, dict] = {}
     if args.adapters:
         print("\n" + "=" * 80)
         print("Processing adapters")
@@ -1022,8 +1042,9 @@ def build():
                     print(msg)
                 discovered_adapters.extend(found)
             elif (path := Path(entry)).is_file() and path.suffix in (".yaml", ".yml"):
-                found = discover_adapters_from_yaml(entry)
+                found, manifest_probe_meta = discover_adapters_from_yaml(entry)
                 discovered_adapters.extend(found)
+                probe_meta.update(manifest_probe_meta)
             else:
                 # Single adapter directory
                 resolved = Path(resolved_path)
@@ -1065,11 +1086,21 @@ def build():
     # token lands, and SR has exactly one usable activation point — the anchor at
     # the end of the generation prompt — so a mislabelled SR checkpoint must not
     # be able to fall through to aLoRA or sequence-start placement.
+    # A probe slot is labelled here for the same reason: its control token is
+    # an end-locator, which is a placement mode, not a weight format, and the
+    # manifest's ``kind`` decides it. An adapter satisfies at most one test -- an SR
+    # adapter carries ``cross_stream`` weights, a probe slot carries a
+    # probe head instead -- and SR goes first because it is read from the
+    # artifact rather than from a label.
     external_discovered = [
         (
             path,
             name,
-            ANCHOR_MODE_SR if path and is_shadow_residual_adapter(path) else tech,
+            ANCHOR_MODE_SR
+            if path and is_shadow_residual_adapter(path)
+            else ANCHOR_MODE_PROBE
+            if probe_meta.get(name, {}).get("kind") == "probe"
+            else tech,
             source,
         )
         for path, name, tech, source in discovered_adapters
@@ -1079,6 +1110,11 @@ def build():
     ]
     all_discovered = external_discovered + built_in_discovered
 
+    # Order LoRA slots first, probe slots last (Python sorts False before True)
+    all_discovered.sort(
+        key=lambda t: probe_meta.get(t[1], {}).get("kind", "lora") == "probe"
+    )
+
     has_external = len(external_discovered) > 0
     has_built_in = len(built_in_discovered) > 0
 
@@ -1227,6 +1263,61 @@ def build():
                 _anchor_text, _sub, _mode = load_activation_anchor(_path, tokenizer)
             adapter_substitute_token_ids[_sub_offset + _i] = _sub
 
+    # adapter_kinds marks which slots are probes (same control token as a
+    # LoRA); resolve each slot's label words to token ids for the vLLM verdict.
+    if probe_meta:
+        adapter_kinds = [
+            probe_meta.get(name, {}).get("kind", "lora") for name in adapter_names
+        ]
+        probe_names = [
+            name
+            for name in adapter_names
+            if probe_meta.get(name, {}).get("kind") == "probe"
+        ]
+
+        if probe_names:
+            # Full width (length num_adapters, indexed by adapter_index - 1): each
+            # probe slot's label token ids. The config derives the per-slot
+            # counts and the padded bank width from this.
+            per_slot_label_names: list[list[str] | None] = []
+            per_slot_label_token_ids: list[list[int] | None] = []
+            # Full width too: each probe slot's read layer (None = final).
+            per_slot_read_layer: list[int | None] = []
+            for name, kind in zip(adapter_names, adapter_kinds):
+                if kind != "probe":
+                    per_slot_label_names.append(None)
+                    per_slot_label_token_ids.append(None)
+                    per_slot_read_layer.append(None)
+                    continue
+                labels = probe_meta[name].get("labels")
+                if not labels:
+                    raise ValueError(
+                        f"Adapter '{name}' is declared kind=probe but has no "
+                        f"'labels' in the manifest. Probe slots need label "
+                        f"words (e.g. labels: [safe, unsafe]) so the verdict exit "
+                        f"can emit them."
+                    )
+                label_names = list(labels)
+                per_slot_label_names.append(label_names)
+                per_slot_label_token_ids.append(
+                    resolve_label_token_ids(tokenizer, label_names)
+                )
+                per_slot_read_layer.append(probe_meta[name].get("layer"))
+
+            optional_kwargs["adapter_kinds"] = adapter_kinds
+            optional_kwargs["probe_label_token_ids"] = per_slot_label_token_ids
+            optional_kwargs["probe_read_layers"] = per_slot_read_layer
+            print("\nProbe slots:")
+            for name, words, ids, layer in zip(
+                adapter_names,
+                per_slot_label_names,
+                per_slot_label_token_ids,
+                per_slot_read_layer,
+            ):
+                if ids is not None:
+                    where = "final norm" if layer is None else f"layer {layer}"
+                    print(f"  {name} → labels {words} → token ids {ids} @ {where}")
+
     model = GraniteSwitchComposer.from_base_and_adapters(
         base_model_name_or_path=base_model_local_path,
         adapter_paths=adapter_paths,