feat: custom zero-shot taxonomies for the v2 detectors

lettucedetect/detectors/llm.py

@@ -117,7 +117,23 @@ def __init__(
         # with their frozen training prompt + hallucinated_spans contract, instead of
         # the zero-shot judge prompt. Auto-detected from the model id; override explicitly.
         self.native = _looks_native(model) if native is None else native
-        if isinstance(include_taxonomy, bool):
+        if self.native and (zero_shot or fewshot_path or prompt_path):
+            raise ValueError(
+                "zero_shot, fewshot_path, and prompt_path have no effect on the native "
+                "generative detectors (they use their frozen training prompt); pass "
+                "native=False to prompt this model as a generic LLM judge instead."
+            )
+        self._custom_taxonomy = not isinstance(include_taxonomy, bool)
+        if (
+            isinstance(include_taxonomy, dict)
+            and include_taxonomy
+            and set(include_taxonomy) <= {"categories", "subcategories"}
+            and all(isinstance(v, dict) for v in include_taxonomy.values())
+        ):
+            # Nested form: full control over both label sets.
+            self.categories = dict(include_taxonomy.get("categories") or {}) or None
+            self.subcategories = dict(include_taxonomy.get("subcategories") or {}) or None
+        elif isinstance(include_taxonomy, bool):
             self.categories = dict(CATEGORY_DEFINITIONS) if include_taxonomy else None
             # The unified taxonomy is category + subcategory; enabling it types both.
             self.subcategories = dict(SUBCATEGORY_DEFINITIONS) if include_taxonomy else None
@@ -424,9 +440,25 @@ def _predict_native(self, context: str, answer: str) -> list[dict]:
         :param answer: The answer string.
         :returns: List of typed spans.
         """
-        system = gen.SYSTEM_EXPL if self.include_reasoning else gen.SYSTEM_BASE
+        if self._custom_taxonomy and self.categories:
+            # Experimental: labels enter the trained prompt only as name/description text,
+            # so a custom taxonomy keeps the structure the model was trained on. Without
+            # custom subcategories the subcategory label collapses to "unspecified".
+            subs = self.subcategories or {
+                "unspecified": gen.SUBCATEGORY_DESCRIPTIONS["unspecified"]
+            }
+            system = gen.build_system_prompt(
+                explain=self.include_reasoning, categories=self.categories, subcategories=subs
+            )
+            schema = build_generative_schema(
+                explain=self.include_reasoning,
+                categories=list(self.categories),
+                subcategories=list(subs),
+            )
+        else:
+            system = gen.SYSTEM_EXPL if self.include_reasoning else gen.SYSTEM_BASE
+            schema = build_generative_schema(explain=self.include_reasoning)
         user = gen.build_user_message(context, answer)
-        schema = build_generative_schema(explain=self.include_reasoning)
 
         cache_key = self.cache._hash(system + user, self.model, str(self.temperature), "native")
         cached = self.cache.get(cache_key)

lettucedetect/detectors/llm_client.py

@@ -103,7 +103,11 @@ def build_hallucination_schema(
     }
 
 
-def build_generative_schema(explain: bool = False) -> dict:
+def build_generative_schema(
+    explain: bool = False,
+    categories: list[str] | None = None,
+    subcategories: list[str] | None = None,
+) -> dict:
     """Build the JSON schema for the fine-tuned span detectors' ``hallucinated_spans`` output.
 
     Distinct from :func:`build_hallucination_schema` (the LLM-judge contract): this
@@ -113,6 +117,8 @@ def build_generative_schema(explain: bool = False) -> dict:
     :mod:`lettucedetect.prompts.generative` label set.
 
     :param explain: Require a per-span ``explanation`` field.
+    :param categories: Replace the frozen category enum (``None`` keeps the default).
+    :param subcategories: Replace the frozen subcategory enum (``None`` keeps the default).
     :returns: JSON schema for ``{"hallucinated_spans": [...]}``.
     """
     from lettucedetect.prompts.generative import (
@@ -122,8 +128,11 @@ def build_generative_schema(explain: bool = False) -> dict:
 
     properties: dict = {
         "text": {"type": "string", "description": "Exact hallucinated substring of the answer"},
-        "category": {"type": "string", "enum": list(CATEGORY_DESCRIPTIONS)},
-        "subcategory": {"type": "string", "enum": list(SUBCATEGORY_DESCRIPTIONS)},
+        "category": {"type": "string", "enum": categories or list(CATEGORY_DESCRIPTIONS)},
+        "subcategory": {
+            "type": "string",
+            "enum": subcategories or list(SUBCATEGORY_DESCRIPTIONS),
+        },
     }
     if explain:
         properties["explanation"] = {

lettucedetect/detectors/taxonomy_head.py

@@ -18,6 +18,31 @@
 _SEP = " [CTX] "  # answer + _SEP + context; must match the taxonomy-head training script
 
 
+def resolve_taxonomy(
+    include_taxonomy: bool | list | dict,
+) -> tuple[dict | None, dict | None]:
+    """Normalize an ``include_taxonomy`` argument into (categories, subcategories) dicts.
+
+    ``True``/``False`` -> frozen defaults (None, None); flat ``{name: description}`` dict ->
+    custom categories; ``{"categories": {...}, "subcategories": {...}}`` -> full control;
+    list of names -> subset of the frozen categories.
+    """
+    if isinstance(include_taxonomy, bool):
+        return None, None
+    if isinstance(include_taxonomy, dict):
+        if (
+            include_taxonomy
+            and set(include_taxonomy) <= {"categories", "subcategories"}
+            and all(isinstance(v, dict) for v in include_taxonomy.values())
+        ):
+            return (
+                dict(include_taxonomy.get("categories") or {}) or None,
+                dict(include_taxonomy.get("subcategories") or {}) or None,
+            )
+        return dict(include_taxonomy) or None, None
+    return {name: CATEGORY_DESCRIPTIONS.get(name) for name in include_taxonomy} or None, None
+
+
 class TaxonomyTyper:
     """Types spans (category + subcategory) via a label-conditioned encoder."""
 
@@ -26,13 +51,19 @@ def __init__(
         model_path: str,
         device: torch.device | str | None = None,
         max_length: int = 1024,
+        categories: dict[str, str] | None = None,
+        subcategories: dict[str, str] | None = None,
         **tok_kwargs: object,
     ) -> None:
         """Load the typing encoder and pre-compute taxonomy-label embeddings.
 
         :param model_path: Path or HF id of the label-conditioned encoder.
         :param device: Device for inference (defaults to CUDA when available).
         :param max_length: Max tokens for the ``answer + context`` input.
+        :param categories: Custom ``{name: description}`` category labels; REPLACES the
+            frozen set (labels enter only as text, so held-out labels type zero-shot from
+            their description). ``None`` keeps the trained taxonomy.
+        :param subcategories: Custom subcategory labels, same semantics.
         :param tok_kwargs: Extra arguments for the tokenizer / model loaders.
         """
         self.device = device or (
@@ -41,15 +72,19 @@ def __init__(
         self.max_length = max_length
         self.tokenizer = AutoTokenizer.from_pretrained(model_path, **tok_kwargs)
         self.encoder = AutoModel.from_pretrained(model_path, **tok_kwargs).to(self.device).eval()
-        self.cat_names = list(CATEGORY_DESCRIPTIONS)
-        self.sub_names = list(SUBCATEGORY_DESCRIPTIONS)
-        self._cat_vecs = self._label_vecs(self.cat_names, CATEGORY_DESCRIPTIONS)
-        self._sub_vecs = self._label_vecs(self.sub_names, SUBCATEGORY_DESCRIPTIONS)
+        cats = categories if categories is not None else CATEGORY_DESCRIPTIONS
+        subs = subcategories if subcategories is not None else SUBCATEGORY_DESCRIPTIONS
+        self.cat_names = list(cats)
+        self.sub_names = list(subs)
+        self._cat_vecs = self._label_vecs(self.cat_names, cats)
+        self._sub_vecs = self._label_vecs(self.sub_names, subs)
 
     @torch.no_grad()
     def _label_vecs(self, names: list[str], descs: dict) -> torch.Tensor:
         enc = self.tokenizer(
-            [f"{n}: {descs[n]}" for n in names], padding=True, return_tensors="pt"
+            [f"{n}: {descs[n]}" if descs.get(n) else n for n in names],
+            padding=True,
+            return_tensors="pt",
         ).to(self.device)
         h = self.encoder(
             input_ids=enc.input_ids, attention_mask=enc.attention_mask

lettucedetect/detectors/transformer.py

@@ -32,6 +32,7 @@ def __init__(
         device: torch.device | str | None = None,
         lang: Lang = "en",
         taxonomy_head: str | None = None,
+        include_taxonomy: bool | list | dict = True,
         **tok_kwargs: object,
     ) -> None:
         """Initialize the transformer detector.
@@ -42,6 +43,11 @@ def __init__(
         :param lang: Language of the model.
         :param taxonomy_head: Optional path/HF id of a label-conditioned typing head. When set,
             ``"spans"`` predictions are typed with a ``category``/``subcategory`` (cascade mode).
+        :param include_taxonomy: Label set for the typing head. ``True`` uses the trained
+            taxonomy; a ``{name: description}`` dict REPLACES the categories (zero-shot: labels
+            enter only as text); ``{"categories": {...}, "subcategories": {...}}`` controls both
+            sets; a list of names selects a subset of the trained categories. Only meaningful
+            together with ``taxonomy_head``.
         :param tok_kwargs: Additional keyword arguments for the tokenizer.
         """
         if lang not in LANG_TO_PASSAGE:
@@ -55,9 +61,15 @@ def __init__(
         self.model.to(self.device).eval()
         self.typer = None
         if taxonomy_head is not None:
-            from lettucedetect.detectors.taxonomy_head import TaxonomyTyper
-
-            self.typer = TaxonomyTyper(taxonomy_head, device=self.device)
+            from lettucedetect.detectors.taxonomy_head import TaxonomyTyper, resolve_taxonomy
+
+            categories, subcategories = resolve_taxonomy(include_taxonomy)
+            self.typer = TaxonomyTyper(
+                taxonomy_head,
+                device=self.device,
+                categories=categories,
+                subcategories=subcategories,
+            )
 
     # ------------------------------------------------------------------
     # Chunking helpers

lettucedetect/prompts/generative.py

@@ -41,10 +41,23 @@
 }
 
 
-def build_system_prompt(explain: bool = False) -> str:
-    """Build the detection system prompt; if ``explain``, also request a per-span explanation."""
-    cats = "\n".join(f"- {k}: {v}" for k, v in CATEGORY_DESCRIPTIONS.items())
-    subs = "\n".join(f"- {k}: {v}" for k, v in SUBCATEGORY_DESCRIPTIONS.items())
+def _label_lines(labels: dict) -> str:
+    return "\n".join(f"- {k}: {v}" if v else f"- {k}" for k, v in labels.items())
+
+
+def build_system_prompt(
+    explain: bool = False,
+    categories: dict[str, str] | None = None,
+    subcategories: dict[str, str] | None = None,
+) -> str:
+    """Build the detection system prompt; if ``explain``, also request a per-span explanation.
+
+    ``categories``/``subcategories`` replace the frozen label sets (labels enter the prompt
+    only as ``name: description`` text, so custom label sets keep the trained structure);
+    ``None`` keeps the frozen defaults.
+    """
+    cats = _label_lines(categories if categories is not None else CATEGORY_DESCRIPTIONS)
+    subs = _label_lines(subcategories if subcategories is not None else SUBCATEGORY_DESCRIPTIONS)
     expl_clause = ", and give a short explanation of why it is unsupported" if explain else ""
     expl_field = ', "explanation": "..."' if explain else ""
     return (