diff --git a/lettucedetect/detectors/llm.py b/lettucedetect/detectors/llm.py index c3f0388..2b7956a 100644 --- a/lettucedetect/detectors/llm.py +++ b/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) diff --git a/lettucedetect/detectors/llm_client.py b/lettucedetect/detectors/llm_client.py index c2cb46f..220d5aa 100644 --- a/lettucedetect/detectors/llm_client.py +++ b/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"] = { diff --git a/lettucedetect/detectors/taxonomy_head.py b/lettucedetect/detectors/taxonomy_head.py index 5851275..7bc7484 100644 --- a/lettucedetect/detectors/taxonomy_head.py +++ b/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,6 +51,8 @@ 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. @@ -33,6 +60,10 @@ def __init__( :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 diff --git a/lettucedetect/detectors/transformer.py b/lettucedetect/detectors/transformer.py index a040496..600046d 100644 --- a/lettucedetect/detectors/transformer.py +++ b/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 diff --git a/lettucedetect/prompts/generative.py b/lettucedetect/prompts/generative.py index cca5f5c..92641d7 100644 --- a/lettucedetect/prompts/generative.py +++ b/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 ( diff --git a/tests/test_custom_taxonomy_pytest.py b/tests/test_custom_taxonomy_pytest.py new file mode 100644 index 0000000..d39d566 --- /dev/null +++ b/tests/test_custom_taxonomy_pytest.py @@ -0,0 +1,217 @@ +"""Tests for custom zero-shot taxonomies across the detector paths (no model, no network).""" + +import json + +import pytest +import torch + +from lettucedetect.detectors.llm import LLMDetector +from lettucedetect.detectors.llm_client import LLMClient, build_generative_schema +from lettucedetect.detectors.taxonomy_head import resolve_taxonomy +from lettucedetect.prompts.generative import build_system_prompt + +CUSTOM_CATS = {"pricing_claim": "any statement about price not supported by the context"} +CUSTOM_SUBS = {"currency": "a currency amount not present in the context"} + + +class TestPromptAndSchemaBuilders: + """Custom labels reach the generative prompt and schema.""" + + def test_system_prompt_defaults_are_frozen(self): + """Test system prompt defaults are frozen.""" + assert "fabricated" in build_system_prompt() or "unsupported" in build_system_prompt() + + def test_system_prompt_custom_categories_replace(self): + """Test system prompt custom categories replace.""" + prompt = build_system_prompt(categories=CUSTOM_CATS, subcategories=CUSTOM_SUBS) + assert "pricing_claim: any statement about price" in prompt + assert "currency: a currency amount" in prompt + assert "unsupported_addition" not in prompt + + def test_system_prompt_label_without_description(self): + """Test system prompt label without description.""" + prompt = build_system_prompt(categories={"bare_label": None}) + assert "- bare_label" in prompt + assert "bare_label: None" not in prompt + + def test_generative_schema_custom_enums(self): + """Test generative schema custom enums.""" + schema = build_generative_schema(categories=["pricing_claim"], subcategories=["currency"]) + item = schema["properties"]["hallucinated_spans"]["items"]["properties"] + assert item["category"]["enum"] == ["pricing_claim"] + assert item["subcategory"]["enum"] == ["currency"] + + def test_generative_schema_defaults_unchanged(self): + """Test generative schema defaults unchanged.""" + item = build_generative_schema()["properties"]["hallucinated_spans"]["items"]["properties"] + assert "unsupported_addition" in item["category"]["enum"] + + +class CaptureClient(LLMClient): + """Stub client capturing the system prompt and schema it is called with.""" + + def __init__(self): + """Init capture store.""" + self.calls = [] + + def complete(self, system, user, model, temperature, schema=None): + """Record the call and return an empty span list.""" + self.calls.append({"system": system, "schema": schema}) + return json.dumps({"hallucinated_spans": []}) + + +class TestNativeCustomTaxonomy: + """Native generative path builds prompt and schema from custom labels.""" + + def make_detector(self, tmp_path, **kwargs): + """Build a native detector with an injected capture client.""" + client = CaptureClient() + det = LLMDetector( + model="KRLabsOrg/lettucedect-v2-qwen-2b", + client=client, + cache_file=str(tmp_path / "cache.json"), + **kwargs, + ) + return det, client + + def test_custom_flat_dict_reaches_prompt_and_schema(self, tmp_path): + """Test custom flat dict reaches prompt and schema.""" + det, client = self.make_detector(tmp_path, include_taxonomy=CUSTOM_CATS) + det.predict(context=["ctx"], answer="ans", question="q", output_format="spans") + call = client.calls[0] + assert "pricing_claim" in call["system"] + item = call["schema"]["properties"]["hallucinated_spans"]["items"]["properties"] + assert item["category"]["enum"] == ["pricing_claim"] + # no custom subcategories: collapses to unspecified + assert item["subcategory"]["enum"] == ["unspecified"] + + def test_nested_dict_controls_both(self, tmp_path): + """Test nested dict controls both.""" + det, client = self.make_detector( + tmp_path, + include_taxonomy={"categories": CUSTOM_CATS, "subcategories": CUSTOM_SUBS}, + ) + det.predict(context=["ctx"], answer="ans", question="q", output_format="spans") + item = client.calls[0]["schema"]["properties"]["hallucinated_spans"]["items"]["properties"] + assert item["subcategory"]["enum"] == ["currency"] + + def test_default_native_prompt_stays_frozen(self, tmp_path): + """Test default native prompt stays frozen.""" + det, client = self.make_detector(tmp_path) + det.predict(context=["ctx"], answer="ans", question="q", output_format="spans") + assert "pricing_claim" not in client.calls[0]["system"] + assert "unsupported_addition" in client.calls[0]["system"] + + def test_native_rejects_dead_args(self, tmp_path): + """Test native rejects dead args.""" + for kwargs in ({"zero_shot": True}, {"fewshot_path": "x.json"}, {"prompt_path": "p.txt"}): + with pytest.raises(ValueError, match="native"): + LLMDetector( + model="KRLabsOrg/lettucedect-v2-qwen-2b", + client=CaptureClient(), + cache_file=str(tmp_path / "c.json"), + **kwargs, + ) + + def test_non_native_still_accepts_zero_shot(self, tmp_path): + """Test non native still accepts zero shot.""" + LLMDetector( + model="gpt-4.1-mini", + zero_shot=True, + client=CaptureClient(), + cache_file=str(tmp_path / "c.json"), + ) + + +class TestResolveTaxonomy: + """resolve_taxonomy normalization.""" + + def test_bool_is_default(self): + """Test bool is default.""" + assert resolve_taxonomy(True) == (None, None) + assert resolve_taxonomy(False) == (None, None) + + def test_flat_dict_is_categories(self): + """Test flat dict is categories.""" + assert resolve_taxonomy(CUSTOM_CATS) == (CUSTOM_CATS, None) + + def test_nested_dict_controls_both(self): + """Test nested dict controls both.""" + cats, subs = resolve_taxonomy({"categories": CUSTOM_CATS, "subcategories": CUSTOM_SUBS}) + assert cats == CUSTOM_CATS and subs == CUSTOM_SUBS + + def test_list_selects_frozen_subset(self): + """Test list selects frozen subset.""" + cats, subs = resolve_taxonomy(["unsupported_addition"]) + assert list(cats) == ["unsupported_addition"] and subs is None + + +class TestTaxonomyTyperCustomLabels: + """Custom labels reach the typer's label-embedding call.""" + + def test_custom_labels_embedded(self, monkeypatch): + """Test custom labels embedded.""" + import lettucedetect.detectors.taxonomy_head as th + + embedded_batches = [] + + class StubEnc(dict): + """Attribute-style batch dict.""" + + def __init__(self, n): + """Build minimal tensors for n inputs.""" + ids = torch.ones((n, 4), dtype=torch.long) + super().__init__(input_ids=ids, attention_mask=torch.ones_like(ids)) + self.input_ids = ids + self.attention_mask = torch.ones_like(ids) + + def to(self, device): + """No-op device move.""" + return self + + class StubTokenizer: + """Records texts passed for embedding.""" + + def __call__(self, texts, **kwargs): + """Tokenize stub.""" + if isinstance(texts, list): + embedded_batches.append(list(texts)) + return StubEnc(len(texts)) + return StubEnc(1) + + class StubOut: + """Encoder output with a deterministic hidden state.""" + + def __init__(self, n): + """Build hidden state.""" + self.last_hidden_state = torch.randn(n, 4, 8) + + class StubModel: + """Minimal encoder stub.""" + + def __call__(self, input_ids=None, attention_mask=None): + """Forward stub.""" + return StubOut(input_ids.shape[0]) + + def to(self, device): + """No-op move.""" + return self + + def eval(self): + """No-op eval.""" + return self + + monkeypatch.setattr( + th, "AutoTokenizer", type("T", (), {"from_pretrained": lambda *a, **k: StubTokenizer()}) + ) + monkeypatch.setattr( + th, "AutoModel", type("M", (), {"from_pretrained": lambda *a, **k: StubModel()}) + ) + typer = th.TaxonomyTyper( + "stub", device="cpu", categories=CUSTOM_CATS, subcategories=CUSTOM_SUBS + ) + assert typer.cat_names == ["pricing_claim"] + assert typer.sub_names == ["currency"] + flat = [x for batch in embedded_batches for x in batch] + assert any("pricing_claim: any statement about price" in x for x in flat) + assert any(x.startswith("currency:") for x in flat)