diff --git a/.gitignore b/.gitignore index cff1c5f2..d3e088de 100644 --- a/.gitignore +++ b/.gitignore @@ -180,4 +180,20 @@ poetry.toml # LSP config files pyrightconfig.json -.vscode \ No newline at end of file +.vscode + +# Claude Code persistent-memory scratch file, not part of the PR +/CLAUDE.md + +# Research/design notes and local PR draft kept on disk for reference, not part of the PR +/docs/archive/ +/docs/research/ +/docs/PR_DESCRIPTION.md + +# Output from scripts/conformal_validation.py +/results/ + +# Leftover scratch files from unrelated branches in this fork +# (feature/vocab-pruning-engine), physically present but irrelevant here +/ROADMAP.md +/pruning_adr.md \ No newline at end of file diff --git a/docs/conformal.md b/docs/conformal.md new file mode 100644 index 00000000..66f9a558 --- /dev/null +++ b/docs/conformal.md @@ -0,0 +1,200 @@ +# Conformal Prediction for GLiNER + +## Stop using `threshold=0.5` + +GLiNER scores every candidate `(span, type)` pair with an independent sigmoid and, by +default, keeps anything above `threshold=0.5`. That number is a convenient default, not +a statistical guarantee — nothing about it tells you what fraction of true entities you're +actually going to miss, and nothing calibrates it to your data, your entity types, or your +risk tolerance. + +`gliner.conformal.ConformalGLiNER` replaces that arbitrary cutoff with a threshold +**calibrated on a held-out labeled set**, backed by finite-sample, distribution-free +guarantees from the conformal prediction literature. Instead of "keep anything above 0.5," +you get to ask for something like: + +> "Calibrate a threshold such that, on average, I miss fewer than 10% of true entities." + +...and get a number back that is provably true (under the assumptions below), not tuned by +eyeballing a validation set. + +This is a new, additive module — it wraps a `GLiNER` model without modifying it, and has +zero effect on the standard `predict_entities`/`inference` API unless you opt in. + +## Quickstart + +```python +from gliner import GLiNER +from gliner.conformal import ConformalGLiNER + +model = GLiNER.from_pretrained("gliner-community/gliner_small-v2.5") +cg = ConformalGLiNER(model) + +# calib_data: held-out labeled sentences, same format as GLiNER's own training/eval data +calib_data = [ + {"tokenized_text": ["Apple", "was", "founded", "by", "Steve", "Jobs", "."], + "ner": [[0, 0, "organization"], [4, 5, "person"]]}, + # ... more labeled examples, ideally 100+ per entity type +] + +cg.calibrate(calib_data, alpha=0.1, mode="risk_control") + +entities = cg.predict_entities( + "Netflix was founded by Reed Hastings.", ["organization", "person"] +) +# each entity carries a "conformal" field: +# {"text": "Netflix", "label": "organization", "score": 0.99, +# "conformal": {"mode": "risk_control", "alpha": 0.1, "calibrated": True}} +``` + +### Shortcut: calibrate on the model itself + +`ConformalGLiNER(model)` above is the full API — `coverage_report`, `save_calibration`, +`thresholds()`, everything. If all you want is to calibrate once and keep predicting from the +same object, `GLiNER` itself exposes a thin convenience wrapper around exactly that: + +```python +model = GLiNER.from_pretrained("gliner-community/gliner_small-v2.5") +model.calibrate(calib_data, alpha=0.1, mode="risk_control") # returns self, chainable + +model.conformal.predict_entities(text, labels) # same ConformalGLiNER instance +model.conformal.thresholds() # per-label thresholds, etc. +``` + +`model.conformal` is `None` until `calibrate()` is called, and is exactly the +`ConformalGLiNER` instance `calibrate()` built — nothing is duplicated between the two APIs, +this just saves constructing the wrapper yourself when you don't need a separate reference to +an uncalibrated model. Same scope restriction applies (`NotImplementedError` on non-span-mode +architectures). + +## The three guarantee modes + +All three are calibrated from the *same* raw span scores GLiNER already computes — no +extra forward pass, no architecture change. + +### `"risk_control"` (the default, and the one to reach for first) + +Bounds the **expected fraction of true entities you miss**, on average across sentences: + +```python +cg.calibrate(calib_data, alpha=0.1, mode="risk_control") +``` + +With `alpha=0.1`: *"on average, ConformalGLiNER misses fewer than 10% of the true entities +in a sentence."* This is the mode to use for compliance/PII-style requirements ("we provably +miss under 5% of PII entities on average") — it directly controls the thing you usually +actually care about (missed entities), rather than an indirect proxy. + +### `"span_filter"` + +Bounds coverage **per entity occurrence**: *"for a randomly drawn true entity of a given +type, it's included in the output with probability at least 90%."* This is the more classical +conformal-prediction framing (closer to "prediction sets" in the broader literature) and is +a good default if you want the simplest possible mental model, or if per-entity behavior +matters more to you than the sentence-level missed-entity rate. + +### `"mondrian"` + +Same as `span_filter`, but calibrated **separately per entity type**, so a common type +(e.g. `person`) can't "subsidize" a rare type (e.g. `chemical_compound`) — each type gets +its own guarantee, at the cost of needing enough calibration examples of *every* type you +care about (see Limitations). + +```python +cg.calibrate(calib_data, alpha=0.1, mode="mondrian") +``` + +## Inspecting the calibrated thresholds + +```python +cg.calibrated_types # ['location', 'organization', 'person'] -- types that got a guarantee +cg.thresholds() # {'location': 0.14, 'organization': 0.09, 'person': 0.11} +``` + +`thresholds()` returns the nonconformity threshold actually applied per label, for every mode: +under `"mondrian"` these genuinely differ per label (that's the point — no type subsidizes +another); under `"span_filter"`/`"risk_control"` every calibrated label currently shares one +pooled value, returned per-label anyway so the API doesn't change shape across modes. Both raise +`RuntimeError` if called before `calibrate()`, same as `predict_entities`/`coverage_report`. + +## Validating and saving a calibration + +```python +report = cg.coverage_report(test_data) # test_data must be disjoint from calib_data +print(report["overall_coverage"], report["per_type_coverage"]) + +cg.save_calibration("calibration.json") +cg2 = ConformalGLiNER.load_calibration("calibration.json", model) +``` + +`test_data` must not overlap with `calib_data` — reusing calibration examples to also +report coverage produces an inflated, meaningless number, since the threshold was tuned +to fit exactly that data. + +## Limitations — read this before you trust a number + +This section exists because a calibrated-looking number is more dangerous than an +obviously-arbitrary one if the calibration doesn't actually apply. + +**The guarantee only covers entity types you actually calibrated on, with enough data.** +Every mode requires roughly `⌈1/alpha⌉` calibration occurrences of a type before it gets a +real threshold (concretely: ~19 for `alpha=0.05`, ~9 for `alpha=0.1`, ~4 for `alpha=0.2`). +If you ask `predict_entities` for a type that wasn't adequately represented in calibration, +`ConformalGLiNER` will: +- warn you loudly, +- fall back to GLiNER's original uncalibrated `p > 0.5` behavior for that type only, +- flag every entity of that type `"conformal": {"calibrated": False}` in the output. + +It will never silently blend an unguaranteed number into a guaranteed-looking one. + +**This is *not* a zero-shot guarantee for arbitrary novel entity types.** This is the most +important limitation and the reason for the point above. Conformal prediction's guarantee +relies on *exchangeability* between your calibration data and what you query at inference +time. If you calibrate on `{person, organization, location}` and then ask for +`chemical_compound` — a type with **zero** calibration occurrences — there is no +mathematical sense in which that query is exchangeable with your calibration set, and no +theorem (here or in the broader conformal-prediction literature) licenses a coverage claim +for it. This isn't a corner case we haven't gotten around to handling; it's a structural +fact about what conformal prediction can prove: split-conformal validity requires the +calibration and test points to be exchangeable, and a type with zero calibration +occurrences was never part of that exchangeable draw at all — there is no rank statistic to +compute a quantile from. GLiNER's flagship feature is arbitrary inference-time label +sets — this module deliberately does *not* pretend to extend a statistical guarantee to +labels outside what you actually calibrated on. If your workflow requires open-vocabulary +guarantees, this isn't (yet) the tool for that; treat the raw sigmoid score as the +heuristic it always was for those types. + +**Domain shift still degrades things, even for calibrated types.** Calibrating on newswire +text and deploying on social media text, for a type name that's nominally the same +(`location` means the same thing in both), is a milder violation of exchangeability than a +genuinely novel type — but it's still a violation. Expect coverage to visibly sag if your +deployment distribution meaningfully differs from your calibration distribution. A concrete +measurement: calibrating on CoNLL-2003 and measuring coverage on WNUT-17, the shared-vocabulary +types (`location`, `person`) still reach 0.87–0.98 coverage across α ∈ {0.05, 0.1, 0.2} — good, +but visibly softer than the ~0.90–0.95 in-domain numbers — while WNUT-17's own types with zero +CoNLL-2003 analogue (`corporation`, `creative-work`, `group`, `product`) sit at a flat ~0.55 +regardless of α, exactly the unguaranteed number you'd expect from a raw uncalibrated cutoff. +Reproduce via `scripts/conformal_validation.py`. + +**`mondrian` mode costs calibration data linearly in the number of types.** Every type +needs its own `~1/alpha`-sized calibration pool; with a fixed calibration budget, more +types means either fewer types getting a real (non-degenerate) threshold, or a looser +`alpha`. + +**Nested-span decoding is applied after the conformal filter, not before.** The coverage/ +risk guarantee is computed against the pre-overlap-resolution candidate set; the final +`predict_entities` output additionally applies GLiNER's usual flat/nested-NER overlap +resolution as a threshold-independent post-processing step. This is a deliberate design +choice needed to keep the Conformal Risk Control guarantee mathematically valid — +applying overlap resolution *before* defining the +calibrated set would break the nesting property the risk-control proof depends on. + +**Scope: span-mode models only.** `ConformalGLiNER` currently supports GLiNER's span-mode +architectures (`UniEncoderSpanGLiNER`, `BiEncoderSpanGLiNER` — the default +`span_mode="markerV0"` configuration, and what most published GLiNER checkpoints use). +Token-mode, generative-decoder, and relation-extraction variants apply their confidence +threshold *inside* the forward pass to prune candidates, so the raw-score interception +this module relies on doesn't give the full candidate universe for those architectures; +using it there would silently understate the true, uncalibrated candidate pool rather than +producing a valid guarantee, so it's explicitly unsupported (raises `NotImplementedError`) +rather than quietly wrong. diff --git a/docs/index.md b/docs/index.md index e2285b26..fb0bf4fa 100644 --- a/docs/index.md +++ b/docs/index.md @@ -17,6 +17,7 @@ architectures add_custom_architectures convert_to_onnx serving +conformal ``` ```{toctree} diff --git a/gliner/conformal/__init__.py b/gliner/conformal/__init__.py new file mode 100644 index 00000000..4e34ab89 --- /dev/null +++ b/gliner/conformal/__init__.py @@ -0,0 +1,20 @@ +"""Conformal-prediction coverage/risk guarantees for GLiNER zero-shot NER. + +See docs/conformal.md for the practitioner-facing guide, including the design +rationale and known limitations. +""" + +from .scores import RawScoreBatch, align_gold_scores, extract_raw_scores +from .wrapper import ConformalGLiNER +from .calibrators import calibration_floor, crc_lambda_search, mondrian_calibrate, split_conformal_quantile + +__all__ = [ + "ConformalGLiNER", + "RawScoreBatch", + "align_gold_scores", + "calibration_floor", + "crc_lambda_search", + "extract_raw_scores", + "mondrian_calibrate", + "split_conformal_quantile", +] diff --git a/gliner/conformal/calibrators.py b/gliner/conformal/calibrators.py new file mode 100644 index 00000000..7c71566b --- /dev/null +++ b/gliner/conformal/calibrators.py @@ -0,0 +1,176 @@ +"""Model-agnostic conformal calibration math. + +Pure NumPy/Python, no GLiNER or PyTorch dependency beyond optional tensor inputs +(anything sequence-like works) -- independently testable against synthetic scores +with analytically known coverage. Implements the three guarantee modes described +in docs/conformal.md: + +- ``split_conformal_quantile``: the ``⌈(n+1)(1-α)⌉``-th order statistic underlying + "span_filter" mode (the standard split-conformal marginal coverage guarantee). +- ``crc_lambda_search``: Conformal Risk Control's finite-sample-conservative + λ search underlying "risk_control" mode. +- ``mondrian_calibrate``: per-type application of ``split_conformal_quantile`` + with an explicit floor, underlying "mondrian" mode. + +All three raise (never silently degrade) when the finite-sample correction has +no solution. +""" + +from __future__ import annotations + +import math +from typing import Dict, Tuple, Mapping, Sequence + + +def calibration_floor(alpha: float) -> int: + """Minimum calibration-set size for which the ``⌈(n+1)(1-α)⌉ ≤ n`` correction is solvable. + + The correction is solvable iff ``n ≥ (1-α)/α``. Returns the smallest integer + n satisfying that. + """ + if not 0 < alpha < 1: + raise ValueError(f"alpha must be in (0, 1), got {alpha}") + return math.ceil((1 - alpha) / alpha) + + +def split_conformal_quantile(scores: Sequence[float], alpha: float) -> float: + """The ``⌈(n+1)(1-α)⌉``-th smallest of ``scores`` -- the split-conformal quantile. + + Args: + scores: Calibration nonconformity scores (larger = worse agreement). + alpha: Miscoverage level in (0, 1). + + Returns: + The conformal quantile ``q̂``; a prediction set ``{y : s(x,y) ≤ q̂}`` then + satisfies ``P(Y ∈ C(X)) ≥ 1-α`` under exchangeability. + + Raises: + ValueError: if ``len(scores) < calibration_floor(alpha)`` -- the quantile + would require a rank beyond the available calibration points + (undefined, not merely wide). + """ + if not 0 < alpha < 1: + raise ValueError(f"alpha must be in (0, 1), got {alpha}") + n = len(scores) + floor = calibration_floor(alpha) + if n < floor: + raise ValueError( + f"n={n} calibration scores insufficient for alpha={alpha}: need n >= {floor} " + f"for the ceil((n+1)(1-alpha))/n correction to be defined. " + "Collect more calibration data or use a larger alpha." + ) + rank = math.ceil((n + 1) * (1 - alpha)) + return sorted(scores)[rank - 1] + + +def mondrian_calibrate( + scores_by_type: Mapping[str, Sequence[float]], alpha: float +) -> Tuple[Dict[str, float], Dict[str, int]]: + """Per-type conformal quantiles, skipping types below the calibration floor. + + Args: + scores_by_type: gold nonconformity scores, grouped by entity type. + alpha: Miscoverage level, shared across all types. + + Returns: + ``(thresholds, skipped)``: ``thresholds`` maps qualifying types to their + per-type quantile; ``skipped`` maps sub-floor types to their observed + calibration count (these fall back to "span_filter"'s pooled threshold + at predict time, not an error here). + """ + thresholds: Dict[str, float] = {} + skipped: Dict[str, int] = {} + for etype, scores in scores_by_type.items(): + try: + thresholds[etype] = split_conformal_quantile(scores, alpha) + except ValueError: + skipped[etype] = len(scores) + return thresholds, skipped + + +def _miss_rate(gold_nc_scores: Sequence[Sequence[float]], lam: float) -> float: + """Mean per-example miss rate ℓ(Cλ,y) at threshold λ (Conformal Risk Control's loss).""" + losses = [] + for example_scores in gold_nc_scores: + if len(example_scores) == 0: + losses.append(0.0) + else: + covered = sum(1 for s in example_scores if s <= lam) + losses.append(1.0 - covered / len(example_scores)) + return sum(losses) / len(losses) if losses else 0.0 + + +def crc_lambda_search( + gold_nc_scores: Sequence[Sequence[float]], + alpha: float, + verify_monotone: bool = True, +) -> float: + """Conformal Risk Control's λ̂ for the missed-entity-rate loss (B=1, bounded in [0,1]). + + ``λ̂ = inf{λ : R̂ₙ(λ) + (1-α)/n ≤ α}``. Candidate λ breakpoints are exactly the + observed nonconformity scores (the loss is a finite step function that only + changes value there), so a grid search over them is exact, not an approximation. + + Args: + gold_nc_scores: one sublist per calibration example, containing + ``1 - p_θ(span,t|x)`` for each of that example's gold entities + (empty sublist for entity-free examples). Use ``float("inf")`` for + gold entities that are structurally unrepresentable (e.g. wider than + ``max_width``) -- they can never be covered, which the loss already + handles correctly without special-casing. + alpha: target expected-miss-rate bound. + verify_monotone: if True, assert the empirical risk is non-increasing + across the candidate grid -- a direct runtime check of the CRC + precondition (GLiNER's independent-sigmoid, single-shared-threshold + decode rule makes the candidate family nested by construction, which + makes this monotone by construction too -- this assertion is a + regression guard on that property, not a hedge against it failing in + practice). Costs one extra pass over the grid; disable only for + large-scale/perf-critical calls after the property has been + established once. + + Returns: + λ̂ ∈ [0, ∞]. ``float("inf")`` means even admitting every candidate + (Cλ = full candidate universe) cannot bring the miss rate to target -- + only possible if some gold entities are structurally unrepresentable in + every example (see the ``float("inf")`` note above). + + Raises: + ValueError: if ``n`` is too small for any λ (including λ=∞) to satisfy + the finite-sample correction: solvable iff ``n ≥ (1-α)/α``, exactly + :func:`calibration_floor` -- the same floor as split conformal, + re-derived independently here from CRC's own formula as a + consistency check. + """ + n = len(gold_nc_scores) + floor = calibration_floor(alpha) + if n < floor: + raise ValueError( + f"n={n} calibration examples insufficient for alpha={alpha}: need n >= {floor} " + "for CRC's finite-sample correction (B-alpha)/n term to be satisfiable even at " + "lambda=infinity. Collect more calibration data or use a larger alpha." + ) + + finite_scores = sorted({s for ex in gold_nc_scores for s in ex if math.isfinite(s)}) + candidates = [0.0, *finite_scores, math.inf] + + rhs = alpha - (1 - alpha) / n + + if verify_monotone: + risks = [_miss_rate(gold_nc_scores, lam) for lam in candidates] + for a, b in zip(risks, risks[1:]): + assert a >= b - 1e-12, ( + "CRC monotonicity precondition violated: empirical risk increased as λ grew. " + "This should be structurally impossible for GLiNER's nested-threshold decode " + "rule -- if this fires, gold_nc_scores was not built from a genuinely nested " + "family of sets." + ) + else: + risks = None + + for i, lam in enumerate(candidates): + risk = risks[i] if risks is not None else _miss_rate(gold_nc_scores, lam) + if risk <= rhs: + return lam + + return math.inf diff --git a/gliner/conformal/scores.py b/gliner/conformal/scores.py new file mode 100644 index 00000000..131cca3d --- /dev/null +++ b/gliner/conformal/scores.py @@ -0,0 +1,142 @@ +"""Raw span-type score extraction for conformal calibration. + +Intercepts GLiNER's forward pass immediately after ``run_batch()``, before +sigmoid/threshold/decode, giving the full dense ``(B, L, K, C)`` candidate +span-score tensor. Reuses ``GLiNER.prepare_base_input`` / +``collate_batch`` / ``run_batch`` directly -- no custom tokenization or collation +logic, no core model changes. + +Scope: span-mode uni-/bi-encoder models only (``UniEncoderSpanGLiNER``, +``BiEncoderSpanGLiNER`` -- the default ``span_mode="markerV0"`` architecture). +Token-mode, decoder, and relex variants apply ``threshold`` *inside* their +forward pass to prune candidate spans before returning scores +(``gliner/modeling/base.py``: ``get_span_representations`` -> +``extract_spans_from_tokens``, and the relex adjacency-selection paths), so +``run_batch()``'s output is not the full candidate universe for those +architectures -- calibrating against it would silently understate true +coverage. Verified by reading every ``forward()`` in ``gliner/modeling/base.py``: +the two span-mode classes never reference ``threshold``, so it is decode-only +for them. +""" + +from __future__ import annotations + +from typing import Any, Dict, List, Tuple, Sequence +from dataclasses import dataclass + +import torch + +_SPAN_MODE_CLASS_NAMES = {"UniEncoderSpanGLiNER", "BiEncoderSpanGLiNER"} + + +def _assert_span_mode_supported(model: Any) -> None: + cls_name = type(model).__name__ + if cls_name not in _SPAN_MODE_CLASS_NAMES: + raise NotImplementedError( + f"ConformalGLiNER v1 only supports span-mode models " + f"({sorted(_SPAN_MODE_CLASS_NAMES)}), got {cls_name!r}. Token-mode, " + "decoder, and relex variants apply `threshold` inside their forward " + "pass to prune candidate spans before returning scores, so " + "run_batch()'s output is not the full candidate universe for those " + "architectures. See docs/conformal.md for details." + ) + + +@dataclass +class RawScoreBatch: + """Raw per-(span,type) scores for one collated batch, pre-sigmoid/threshold/decode. + + Attributes: + logits: ``(B, L, K, C)`` raw span-mode scores, pre-sigmoid. + id_to_classes: per-item ``{1-indexed class id: type string}`` maps + (0 is reserved/unused, matching ``gliner/decoding/decoder.py``'s convention). + tokens: per-item word-token lists, aligned with the ``(start, end)`` + word indices in each example's gold ``ner`` triples. + """ + + logits: torch.Tensor + id_to_classes: List[Dict[int, str]] + tokens: List[List[str]] + + +def extract_raw_scores(model: Any, examples: Sequence[Dict[str, Any]], labels: Sequence[str]) -> RawScoreBatch: + """Run one forward pass and return dense pre-sigmoid span-type scores. + + Args: + model: A span-mode ``GLiNER`` instance. + examples: Pre-tokenized examples, ``{"tokenized_text": List[str], "ner": ...}`` + (the ``"ner"`` field is ignored here; use :func:`align_gold_scores` to pull + out gold-span scores). Passing already-tokenized words (rather than raw + text through ``model.prepare_batch``) is deliberate: it guarantees the + word indices in ``examples[i]["ner"]`` line up exactly with the model's + own span indexing, with no re-tokenization drift. + labels: The fixed target label set to score every example against. + + Returns: + RawScoreBatch with the dense score tensor and per-item bookkeeping. + """ + _assert_span_mode_supported(model) + if not examples: + raise ValueError("No examples to score.") + + all_tokens = [ex["tokenized_text"] for ex in examples] + input_x = model.prepare_base_input(all_tokens) + batch = model.collate_batch(input_x, list(labels)) + model_output = model.run_batch(batch, threshold=0.0, move_to_device=True) + + logits = model_output.logits if hasattr(model_output, "logits") else model_output[0] + if not isinstance(logits, torch.Tensor): + logits = torch.from_numpy(logits) + + id_to_classes = batch["id_to_classes"] + if not isinstance(id_to_classes, list): + id_to_classes = [id_to_classes] * logits.shape[0] + + return RawScoreBatch(logits=logits, id_to_classes=id_to_classes, tokens=batch["tokens"]) + + +def align_gold_scores( + raw: RawScoreBatch, + examples: Sequence[Dict[str, Any]], +) -> Tuple[List[float], List[str], List[int]]: + """Pull nonconformity scores ``1 - sigmoid(logit)`` for every gold ``(span, type)`` pair. + + Args: + raw: Output of :func:`extract_raw_scores` for the same ``examples``. + examples: Same list passed to :func:`extract_raw_scores` (must match order/length). + + Returns: + Tuple of parallel lists ``(scores, types, example_idx)``: nonconformity score, + gold entity type, and the index into ``examples`` it came from. A gold span + wider than ``max_width`` (not representable in the candidate universe at all -- + GLiNER structurally cannot ever predict it) gets score ``float("inf")`` -- + guaranteed non-conforming, guaranteed "missed" under + risk-control, exactly the correct behavior for an unrepresentable entity, not + a special case to filter out. + """ + if len(examples) != len(raw.id_to_classes): + raise ValueError(f"examples/raw batch size mismatch: {len(examples)} vs {len(raw.id_to_classes)}") + + probs = torch.sigmoid(raw.logits) + _, L, K, C = probs.shape + + scores: List[float] = [] + types: List[str] = [] + example_idx: List[int] = [] + + for i, ex in enumerate(examples): + class_to_id = {v: k for k, v in raw.id_to_classes[i].items()} + for start, end, etype in ex.get("ner", []): + width_offset = end - start + if etype not in class_to_id: + continue # type not in this batch's label set -- not calibratable from this call + col = class_to_id[etype] - 1 # id_to_classes is 1-indexed (0 reserved) + if not (0 <= start < L) or not (0 <= width_offset < K) or not (0 <= col < C): + score = float("inf") + else: + score = 1.0 - probs[i, start, width_offset, col].item() + scores.append(score) + types.append(etype) + example_idx.append(i) + + return scores, types, example_idx diff --git a/gliner/conformal/wrapper.py b/gliner/conformal/wrapper.py new file mode 100644 index 00000000..d3c8838d --- /dev/null +++ b/gliner/conformal/wrapper.py @@ -0,0 +1,445 @@ +"""ConformalGLiNER -- conformal-prediction wrapper around a span-mode GLiNER model. + +See docs/conformal.md for the design rationale. Summary of the one behavior +every method below enforces: the ``>= 1-alpha`` guarantee applies only to +entity types adequately represented in the calibration set (``>= +calibration_floor(alpha)`` gold occurrences). Any other type is served from +GLiNER's original uncalibrated ``p > 0.5`` rule, flagged ``"calibrated": False``, +with a loud warning -- never silently blended into a guaranteed-looking number. +""" + +from __future__ import annotations + +import json +import warnings +from typing import Any, Dict, List, Union, Optional, Sequence +from collections import Counter, defaultdict +from dataclasses import field, dataclass + +import torch + +from gliner.decoding.decoder import Span + +from .scores import align_gold_scores, extract_raw_scores +from .calibrators import calibration_floor, crc_lambda_search, mondrian_calibrate, split_conformal_quantile + +_VALID_MODES = {"span_filter", "risk_control", "mondrian"} + + +@dataclass +class _CalibrationState: + mode: str + alpha: float + labels: List[str] + calibrated_types: List[str] + type_counts: Dict[str, int] + pooled_nc_threshold: Optional[float] = None + mondrian_thresholds: Dict[str, float] = field(default_factory=dict) + mondrian_skipped: Dict[str, int] = field(default_factory=dict) + crc_lambda: Optional[float] = None + model_id: Optional[str] = None + + def to_json_dict(self) -> Dict[str, Any]: + return { + "mode": self.mode, + "alpha": self.alpha, + "labels": self.labels, + "calibrated_types": self.calibrated_types, + "type_counts": self.type_counts, + "pooled_nc_threshold": self.pooled_nc_threshold, + "mondrian_thresholds": self.mondrian_thresholds, + "mondrian_skipped": self.mondrian_skipped, + "crc_lambda": self.crc_lambda, + "model_id": self.model_id, + } + + @classmethod + def from_json_dict(cls, d: Dict[str, Any]) -> _CalibrationState: + return cls(**d) + + +class ConformalGLiNER: + """Wraps a span-mode GLiNER model with a calibrated conformal filter. + + Never mutates the wrapped model. See docs/conformal.md for the API rationale + and exactly what the guarantee does and does not cover. + """ + + def __init__(self, model: Any): + self.model = model + self._state: Optional[_CalibrationState] = None + + @property + def is_calibrated(self) -> bool: + return self._state is not None + + @property + def calibrated_types(self) -> List[str]: + """Entity types that reached the calibration floor and carry the coverage guarantee. + + Any label requested at predict time that is *not* in this list falls back to + GLiNER's original uncalibrated ``p > 0.5`` rule -- see ``predict_entities``. + """ + return list(self._require_calibrated().calibrated_types) + + def thresholds(self) -> Dict[str, float]: + """Per-label nonconformity threshold actually applied at prediction time. + + Keyed by every entry in ``calibrated_types``. For ``mode="mondrian"`` these + differ per label by design (that's the whole point of Mondrian calibration -- + no type "subsidizes" another). For ``"span_filter"``/``"risk_control"`` every + calibrated label currently shares one pooled threshold/lambda; returned + per-label here anyway for a uniform API across modes, not because the value + differs. Raises if not yet calibrated, same as every other query method here. + """ + state = self._require_calibrated() + return {etype: self._nc_threshold_for(state, etype) for etype in state.calibrated_types} + + def _require_calibrated(self) -> _CalibrationState: + if self._state is None: + raise RuntimeError("ConformalGLiNER is not calibrated. Call calibrate() first.") + return self._state + + @staticmethod + def _model_id(model: Any) -> Optional[str]: + return getattr(getattr(model, "config", None), "_name_or_path", None) + + # ------------------------------------------------------------------ # + # Calibration + # ------------------------------------------------------------------ # + + def calibrate( + self, + calib_data: Sequence[Dict[str, Any]], + alpha: float, + mode: str = "risk_control", + labels: Optional[Sequence[str]] = None, + ) -> ConformalGLiNER: + """Calibrate the conformal threshold(s) on a held-out labeled set. + + Args: + calib_data: ``[{"tokenized_text": [...], "ner": [[start,end,type],...]}, ...]`` + -- the same schema GLiNER's own training/eval pipeline uses + (gliner/data_processing/processor.py). Must be disjoint from any + data later passed to :meth:`coverage_report` -- reusing + calibration examples to also measure coverage produces a + biased, inflated estimate. + alpha: target miscoverage/risk level in (0, 1). + mode: one of ``"span_filter"``, ``"risk_control"``, ``"mondrian"``. + No default is silently assumed by the public API surface beyond + this parameter's own default; callers relying on the default + should be aware it is ``"risk_control"``. + labels: the fixed target label set 𝒯_cal. Defaults to every type + appearing at least once in ``calib_data``. + + Returns: + ``self``, for chaining. + + Raises: + ValueError: invalid ``mode``/``alpha``, or too few calibration + examples for the requested ``alpha`` -- raises rather than + silently degrading. + """ + if mode not in _VALID_MODES: + raise ValueError(f"mode must be one of {sorted(_VALID_MODES)}, got {mode!r}") + if not 0 < alpha < 1: + raise ValueError(f"alpha must be in (0, 1), got {alpha}") + if not calib_data: + raise ValueError("calib_data is empty.") + + if labels is None: + labels = sorted({etype for ex in calib_data for (_, _, etype) in ex.get("ner", [])}) + labels = list(labels) + if not labels: + raise ValueError("No labels found in calib_data and none provided explicitly.") + + raw = extract_raw_scores(self.model, calib_data, labels) + scores, types, example_idx = align_gold_scores(raw, calib_data) + + floor = calibration_floor(alpha) + type_counts = Counter(types) + calibrated_types = sorted(t for t, n in type_counts.items() if n >= floor) + if not calibrated_types: + raise ValueError( + f"No requested type reached the calibration floor (>= {floor} gold occurrences " + f"needed for alpha={alpha}). Observed counts: {dict(type_counts)}. Collect more " + "calibration data, request fewer/more-common types, or use a larger alpha." + ) + under_floor = {t: n for t, n in type_counts.items() if n < floor} + if under_floor: + warnings.warn( + f"Type(s) {under_floor} have fewer than {floor} gold calibration occurrences " + f"(alpha={alpha}) and will NOT receive a calibrated guarantee at predict time " + "(raw uncalibrated p>0.5 fallback will be used for them, flagged accordingly).", + UserWarning, + stacklevel=2, + ) + + state = _CalibrationState( + mode=mode, + alpha=alpha, + labels=labels, + calibrated_types=calibrated_types, + type_counts=dict(type_counts), + model_id=self._model_id(self.model), + ) + + if mode in ("span_filter", "mondrian"): + # Pooled threshold: the marginal-over-calibrated-types guarantee, and (for + # mondrian) the fallback for any calibrated-but-not-enough-for-its-own- + # Mondrian-cell type -- though by construction every type in + # `calibrated_types` already met the same floor, so mondrian_calibrate below + # should not skip any of them; the pooled value is kept regardless as the + # documented, deterministic fallback path. + pooled_scores = [s for s, t in zip(scores, types) if t in calibrated_types] + state.pooled_nc_threshold = split_conformal_quantile(pooled_scores, alpha) + + if mode == "mondrian": + scores_by_type: Dict[str, List[float]] = defaultdict(list) + for s, t in zip(scores, types): + if t in calibrated_types: + scores_by_type[t].append(s) + state.mondrian_thresholds, state.mondrian_skipped = mondrian_calibrate(scores_by_type, alpha) + + if mode == "risk_control": + gold_nc_scores: List[List[float]] = [[] for _ in calib_data] + for s, t, i in zip(scores, types, example_idx): + if t in calibrated_types: + gold_nc_scores[i].append(s) + state.crc_lambda = crc_lambda_search(gold_nc_scores, alpha) + + self._state = state + return self + + # ------------------------------------------------------------------ # + # Inference + # ------------------------------------------------------------------ # + + def _nc_threshold_for(self, state: _CalibrationState, etype: str) -> float: + if state.mode == "risk_control": + return state.crc_lambda + if state.mode == "span_filter": + return state.pooled_nc_threshold + if state.mode == "mondrian": + return state.mondrian_thresholds.get(etype, state.pooled_nc_threshold) + raise AssertionError(f"unreachable mode {state.mode!r}") + + def predict_entities( + self, + text: Union[str, List[str]], + labels: Sequence[str], + flat_ner: bool = True, + multi_label: bool = False, + ) -> Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]]: + """Predict entities with conformal-guarantee metadata attached. + + Behaves like ``GLiNER.predict_entities``/``batch_predict_entities`` in + shape (single text -> flat list; list of texts -> list of lists), but + the admission rule is the calibrated conformal threshold, not a raw 0.5 + cutoff, for every type in ``labels`` that was adequately represented at + calibration time. Every returned entity carries a + ``"conformal": {"mode", "alpha", "calibrated"}`` field; + ``"calibrated": False`` means that entity's type had no valid + guarantee and was produced by the original uncalibrated rule instead. + """ + state = self._require_calibrated() + single = isinstance(text, str) + texts = [text] if single else list(text) + labels = list(labels) + + prepared = self.model.prepare_batch(texts, labels) + if not prepared["valid_texts"]: + empty: List[List[Dict[str, Any]]] = [[] for _ in texts] + return empty[0] if single else empty + + batch = self.model.collate_batch(prepared["input_x"], prepared["entity_types"]) + model_output = self.model.run_batch(batch, threshold=0.0, move_to_device=True) + logits = model_output.logits if hasattr(model_output, "logits") else model_output[0] + if not isinstance(logits, torch.Tensor): + logits = torch.from_numpy(logits) + probs = torch.sigmoid(logits) + B, _, _, C = probs.shape + + id_to_classes = batch["id_to_classes"] + if not isinstance(id_to_classes, list): + id_to_classes = [id_to_classes] * B + + num_tokens = [len(t) for t in batch["tokens"]] + + uncalibrated_requested: set = set() + decoded_per_item: List[List[Any]] = [] + + for b in range(B): + cls_map = id_to_classes[b] + spans: List[Span] = [] + for col in range(C): + etype = cls_map.get(col + 1) + if etype is None: + continue + calibrated = etype in state.calibrated_types + if not calibrated: + uncalibrated_requested.add(etype) + admit_col = probs[b, :, :, col] > 0.5 + else: + tau = self._nc_threshold_for(state, etype) + admit_col = (1.0 - probs[b, :, :, col]) <= tau + s_idx, k_idx = torch.where(admit_col) + for s, k in zip(s_idx.tolist(), k_idx.tolist()): + if s + k >= num_tokens[b]: + continue + score = probs[b, s, k, col].item() + spans.append(Span(start=s, end=s + k, entity_type=etype, score=score)) + decoded_per_item.append(self.model.decoder.greedy_search(spans, flat_ner=flat_ner, multi_label=multi_label)) + + if uncalibrated_requested: + warnings.warn( + f"Type(s) {sorted(uncalibrated_requested)} were not adequately represented in " + f"calibration and have NO conformal guarantee -- served via GLiNER's original " + "uncalibrated p>0.5 rule instead. Entities of these types are flagged " + '"conformal": {"calibrated": False} in the output.', + UserWarning, + stacklevel=2, + ) + + entities = self.model.map_entities_to_text( + decoded_per_item, + prepared["valid_texts"], + prepared["valid_to_orig_idx"], + prepared["start_token_map"], + prepared["end_token_map"], + prepared["num_original"], + ) + for per_text in entities: + for ent in per_text: + calibrated = ent["label"] in state.calibrated_types + ent["conformal"] = {"mode": state.mode, "alpha": state.alpha, "calibrated": calibrated} + + return entities[0] if single else entities + + # ------------------------------------------------------------------ # + # Empirical validation + # ------------------------------------------------------------------ # + + def coverage_report( + self, test_data: Sequence[Dict[str, Any]], labels: Optional[Sequence[str]] = None + ) -> Dict[str, Any]: + """Empirically measure coverage/efficiency on held-out labeled data. + + ``test_data`` must be disjoint from whatever was passed to + :meth:`calibrate` -- reusing calibration data here trivially inflates + the coverage estimate, since the threshold was tuned to fit exactly + that data. This method does not enforce disjointness itself (it has no + way to know the calibration set's identity at this layer); callers/ + tests are responsible. + + Returns a dict with overall + per-type coverage, restricted to + calibrated types -- never blended with uncalibrated ones -- and + efficiency (mean admitted candidates per example). + + ``overall_coverage`` reports the quantity actually calibrated for + ``state.mode``, not a one-size-fits-all pooled statistic: for + ``"span_filter"``/``"mondrian"`` that's the marginal per-entity coverage + (pooled over every gold entity); for ``"risk_control"`` it's + ``1 - mean_per_sentence_miss_rate``, matching Conformal Risk Control's + own loss definition exactly. These are genuinely different quantities + whenever gold-entity count varies across sentences -- pooling entities + flat for risk_control would silently report an uncalibrated number and + can show spurious undercoverage unrelated to whether the actual CRC + guarantee holds. (Caught empirically while validating this module.) + """ + state = self._require_calibrated() + labels = list(labels) if labels else list(state.labels) + + raw = extract_raw_scores(self.model, test_data, labels) + scores, types, example_idx = align_gold_scores(raw, test_data) + + per_type_hits: Dict[str, int] = defaultdict(int) + per_type_n: Dict[str, int] = defaultdict(int) + n_uncalibrated_gold = 0 + per_example_gold: Dict[int, List[bool]] = defaultdict(list) + for s, t, ex_i in zip(scores, types, example_idx): + if t not in state.calibrated_types: + n_uncalibrated_gold += 1 + continue + tau = self._nc_threshold_for(state, t) + hit = s <= tau + per_type_n[t] += 1 + per_type_hits[t] += int(hit) + per_example_gold[ex_i].append(hit) + + total_n = sum(per_type_n.values()) + total_hits = sum(per_type_hits.values()) + + if state.mode == "risk_control": + sentence_losses = [ + 1.0 - sum(hits) / len(hits) if hits else 0.0 + for hits in (per_example_gold.get(i, []) for i in range(len(test_data))) + ] + overall_coverage = 1.0 - sum(sentence_losses) / len(sentence_losses) if sentence_losses else float("nan") + else: + overall_coverage = (total_hits / total_n) if total_n else float("nan") + + # Efficiency: mean admitted (span,type) pairs per example, over the full dense + # candidate grid (not just gold cells) -- reuses the same forward pass, no extra cost. + probs = torch.sigmoid(raw.logits) + B = probs.shape[0] + admitted_counts = torch.zeros(B) + raw_candidate_counts = torch.zeros(B) + for b in range(B): + cls_map = raw.id_to_classes[b] + for col in range(probs.shape[3]): + etype = cls_map.get(col + 1) + if etype is None or etype not in state.calibrated_types: + continue + tau = self._nc_threshold_for(state, etype) + admitted_counts[b] += ((1.0 - probs[b, :, :, col]) <= tau).sum().item() + raw_candidate_counts[b] += probs.shape[1] * probs.shape[2] + + return { + "mode": state.mode, + "alpha": state.alpha, + "n_test_examples": len(test_data), + "overall_coverage": overall_coverage, + "n_calibrated_gold": total_n, + "n_uncalibrated_gold": n_uncalibrated_gold, + "per_type_coverage": {t: per_type_hits[t] / per_type_n[t] for t in per_type_n}, + "per_type_n": dict(per_type_n), + "efficiency_mean": admitted_counts.mean().item(), + "raw_candidates_mean": raw_candidate_counts.mean().item(), + } + + # ------------------------------------------------------------------ # + # Serialization + # ------------------------------------------------------------------ # + + def save_calibration(self, path: str) -> None: + """Serialize calibration state (not the model) to JSON.""" + state = self._require_calibrated() + with open(path, "w") as f: + json.dump(state.to_json_dict(), f, indent=2) + + @classmethod + def load_calibration(cls, path: str, model: Any) -> ConformalGLiNER: + """Re-wrap ``model`` with a previously saved calibration state. + + Warns (does not raise) if ``model``'s identity doesn't match the model + the calibration was computed against -- nonconformity scores are + model-specific, so a mismatch means the loaded thresholds may not carry + a valid guarantee for this model, but a deliberate same-architecture + swap (e.g. a re-exported checkpoint) is a legitimate use case. + """ + with open(path) as f: + d = json.load(f) + state = _CalibrationState.from_json_dict(d) + current_id = cls._model_id(model) + if state.model_id is not None and current_id is not None and state.model_id != current_id: + warnings.warn( + f"Loaded calibration was computed against model {state.model_id!r}, but this " + f"model is {current_id!r}. Nonconformity scores are model-specific -- the " + "guarantee may not hold unless this is a deliberate, compatible swap.", + UserWarning, + stacklevel=2, + ) + cg = cls(model) + cg._state = state + return cg diff --git a/gliner/model.py b/gliner/model.py index c666fa9b..0afd5a30 100644 --- a/gliner/model.py +++ b/gliner/model.py @@ -2413,6 +2413,51 @@ def batch_predict_entities( **kwargs, ) + def calibrate( + self, + calib_data: List[Dict[str, Any]], + alpha: float, + mode: str = "risk_control", + labels: Optional[List[str]] = None, + ) -> "BaseEncoderGLiNER": + """Calibrate this model with a conformal coverage/risk guarantee. + + Thin convenience wrapper around `gliner.conformal.ConformalGLiNER`: builds + one around `self`, calibrates it, and stores it on this instance (accessible + via the `conformal` property) so calibration and inference live on the same + object instead of requiring callers to juggle a separate wrapper. For + anything beyond predicting with the calibrated threshold -- coverage_report, + save_calibration/load_calibration, inspecting per-label thresholds -- use + `self.conformal` directly, or construct `ConformalGLiNER(model)` yourself; + this method does not duplicate that surface. + + Only span-mode models (the default `UniEncoderSpanGLiNER`/`BiEncoderSpanGLiNER` + architecture) are supported; other architectures raise `NotImplementedError` + from `ConformalGLiNER` itself -- see `docs/conformal.md`'s Scope section. + + Args: + calib_data: Held-out labeled examples, same schema as GLiNER's own + training/eval data: `[{"tokenized_text": [...], "ner": [[start, + end, type], ...]}, ...]`. Must be disjoint from any data later + passed to `self.conformal.coverage_report(...)`. + alpha: Target miscoverage/risk level in (0, 1). + mode: One of `"span_filter"`, `"risk_control"`, `"mondrian"`. + labels: The fixed target label set to calibrate against. Defaults to + every type appearing at least once in `calib_data`. + + Returns: + `self`, for chaining (e.g. `model.calibrate(data, alpha=0.1).predict_entities(...)`). + """ + from .conformal import ConformalGLiNER # noqa: PLC0415 (opt-in, not a hard dependency) + + self._conformal_model = ConformalGLiNER(self).calibrate(calib_data, alpha=alpha, mode=mode, labels=labels) + return self + + @property + def conformal(self) -> Optional[Any]: + """The `ConformalGLiNER` wrapper built by `calibrate()`, or `None` if not yet calibrated.""" + return getattr(self, "_conformal_model", None) + @torch.no_grad() def evaluate( self, diff --git a/pyproject.toml b/pyproject.toml index e22ec6ff..bea7f20f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -171,6 +171,8 @@ ignore = [ "RUF012", # Too many arguments "PLR0913", + # Too many positional arguments + "PLR0917", # Too many branches "PLR0912", # Too many statements diff --git a/scripts/conformal_validation.py b/scripts/conformal_validation.py new file mode 100644 index 00000000..8e053af9 --- /dev/null +++ b/scripts/conformal_validation.py @@ -0,0 +1,562 @@ +"""Empirical validation of ConformalGLiNER's coverage/risk guarantees. + +Validates against real data (not synthetic): CoNLL-2003 and WNUT-17 via +DFKI-SLT/cross_ner (sidesteps `datasets`'s script-loading rejection for these +two datasets), using gliner-community/gliner_small-v2.5. + +Scope disclosed up front, not hidden: this run covers in-domain CoNLL-2003, +in-domain WNUT-17, and a zero-shot pair (calibrate on CoNLL-2003, measure +coverage on WNUT-17) -- not a full CrossNER multi-domain sweep. Pool sizes +are capped (see POOL_CAP below) for CPU runtime; T defaults to 50 trials so +this is runnable in one sitting on a laptop (bump for final/published +numbers). Both are disclosed in the output results markdown too. + +One forward pass per pooled sentence set; all T-trial resampling happens +on cached scores/tensors afterward (no repeated model calls per trial). + +Usage: + OMP_NUM_THREADS=1 KMP_DUPLICATE_LIB_OK=TRUE python scripts/conformal_validation.py \ + --output_dir results/conformal +""" + +from __future__ import annotations + +import os +import json +import time +import random +import argparse +from typing import Dict, List, Tuple, Sequence +from pathlib import Path +from dataclasses import dataclass + +os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE") +os.environ.setdefault("OMP_NUM_THREADS", "1") +os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") + +import matplotlib + +matplotlib.use("Agg") +import torch +import matplotlib.pyplot as plt +from datasets import load_dataset + +from gliner import GLiNER +from gliner.conformal.scores import extract_raw_scores +from gliner.conformal.calibrators import calibration_floor, crc_lambda_search, split_conformal_quantile + +MODEL_ID = "gliner-community/gliner_small-v2.5" +ALPHAS = [0.05, 0.10, 0.20] +POOL_CAP = 1200 # sentences per pool -- see module docstring +BATCH_SIZE = 16 + + +def bio_to_spans(tags: List[str]) -> List[Tuple[int, int, str]]: + """Decode a BIO tag sequence into inclusive-end (start, end, type) triples.""" + spans = [] + start = None + etype = None + for i, tag in enumerate([*tags, "O"]): + if tag.startswith("B-"): + if start is not None: + spans.append((start, i - 1, etype)) + start, etype = i, tag[2:] + elif tag.startswith("I-") and etype == tag[2:]: + continue + else: + if start is not None: + spans.append((start, i - 1, etype)) + start, etype = None, None + return spans + + +def load_examples(dataset_id: str, config: str, split: str, cap: int) -> List[Dict]: + ds = load_dataset(dataset_id, name=config, split=split) if config else load_dataset(dataset_id, split=split) + tag_names = ds.features["ner_tags"].feature.names + out = [] + for row in ds: + if cap is not None and len(out) >= cap: + break + tokens = row["tokens"] + if not tokens: + continue + tags = [tag_names[t] for t in row["ner_tags"]] + ner = [list(s) for s in bio_to_spans(tags)] + out.append({"tokenized_text": tokens, "ner": ner}) + return out + + +@dataclass +class Pool: + name: str + examples: List[Dict] + probs: List[torch.Tensor] # per-example (L, K, C) sigmoid probs + id_to_class: List[Dict[int, str]] + labels: List[str] + + +def build_pool(model, name: str, examples: List[Dict], labels: Sequence[str], batch_size: int = BATCH_SIZE) -> Pool: + probs: List[torch.Tensor] = [] + id_to_class: List[Dict[int, str]] = [] + for i in range(0, len(examples), batch_size): + batch = examples[i : i + batch_size] + raw = extract_raw_scores(model, batch, labels) + p = torch.sigmoid(raw.logits) + for j in range(p.shape[0]): + probs.append(p[j]) + id_to_class.append(raw.id_to_classes[j]) + return Pool(name=name, examples=examples, probs=probs, id_to_class=id_to_class, labels=list(labels)) + + +def gold_nc_by_example(pool: Pool) -> List[List[Tuple[str, float]]]: + """Per example, list of (type, nonconformity_score) for gold entities. + + Unrepresentable gold spans (wider than max_width) get score=inf, see + gliner/conformal/scores.py::align_gold_scores -- reimplemented here + per-example since Pool caches dense per-example tensors rather than a + flat batch. + """ + out = [] + for ex, probs, cls_map in zip(pool.examples, pool.probs, pool.id_to_class): + class_to_id = {v: k for k, v in cls_map.items()} + L, K, C = probs.shape + entry = [] + for start, end, etype in ex.get("ner", []): + if etype not in class_to_id: + continue + width = end - start + col = class_to_id[etype] - 1 + if not (0 <= start < L) or not (0 <= width < K) or not (0 <= col < C): + score = float("inf") + else: + score = 1.0 - probs[start, width, col].item() + entry.append((etype, score)) + out.append(entry) + return out + + +def trial_metrics( + calib_pool: Pool, + test_pool: Pool, + calib_gold: List[List[Tuple[str, float]]], + test_gold: List[List[Tuple[str, float]]], + alpha: float, + n_calib: int, + n_trials: int, + mode: str, + seed: int, + pool_and_resplit: bool = False, +) -> Dict: + """Mirror ConformalGLiNER's own calibrated/uncalibrated split. + + A type only contributes to the headline coverage/efficiency numbers if it met + the calibration floor in *that trial's* calibration subsample. Types requested + at test time that never met the floor (e.g. WNUT-only types under Pair A's + CoNLL-derived calibration) are tracked separately as `uncalibrated_*` -- never + blended into the guaranteed-looking headline number. This is exactly the + scenario the zero-shot descope predicts and this eval is meant to + demonstrate, not accidentally paper over. + + pool_and_resplit=True implements the correct in-domain protocol: pool + calib_pool+test_pool together and draw a *fresh* random calib/test + partition every trial, rather than using calib_pool and test_pool as static, + separately-sourced sets. This matters empirically, not just by-the-book: an + earlier run of this script found CoNLL-2003's *official* validation and test + splits are themselves not fully exchangeable for this model (mean + nonconformity 0.22 on validation vs 0.27 on test -- a real, documented + property of that benchmark's val/test construction, not a code bug), which + silently violated split conformal's exchangeability precondition and produced + a measured ~4-5pp coverage undershoot. Pooling and re-splitting per trial is + the correct way to test "does split-conformal coverage hold when + exchangeability genuinely is satisfied" without that confound. Pair A + (zero-shot) deliberately keeps calib_pool/test_pool separate -- that + non-exchangeability *is* the experiment there. + """ + rng = random.Random(seed) + + if pool_and_resplit: + combined_probs = calib_pool.probs + test_pool.probs + combined_id_to_class = calib_pool.id_to_class + test_pool.id_to_class + combined_gold = calib_gold + test_gold + combined_n = len(combined_probs) + else: + calib_n_total = len(calib_pool.examples) + test_indices_all = list(range(len(test_pool.examples))) + + coverages, effs, raw_counts = [], [], [] + uncal_coverages = [] + per_type_hits: Dict[str, int] = {} + per_type_n: Dict[str, int] = {} + n_ok_trials = 0 + floor = calibration_floor(alpha) + + for _trial in range(n_trials): + if pool_and_resplit: + shuffled = list(range(combined_n)) + rng.shuffle(shuffled) + calib_idx = shuffled[: min(n_calib, combined_n)] + test_idx = shuffled[min(n_calib, combined_n) :] + else: + calib_idx = rng.sample(range(calib_n_total), min(n_calib, calib_n_total)) + test_idx = test_indices_all + + calib_types_n: Dict[str, int] = {} + for i in calib_idx: + for t, _ in (combined_gold if pool_and_resplit else calib_gold)[i]: + calib_types_n[t] = calib_types_n.get(t, 0) + 1 + calibrated_types = {t for t, n in calib_types_n.items() if n >= floor} + if not calibrated_types: + continue + calib_gold_source = combined_gold if pool_and_resplit else calib_gold + pooled_scores = [s for i in calib_idx for (t, s) in calib_gold_source[i] if t in calibrated_types] + if len(pooled_scores) < floor: + continue + + if mode == "span_filter": + try: + tau = split_conformal_quantile(pooled_scores, alpha) + except ValueError: + continue + + def admit(s, tau=tau): + return s <= tau + else: # risk_control + gold_lists = [[s for (t, s) in calib_gold_source[i] if t in calibrated_types] for i in calib_idx] + try: + lam = crc_lambda_search(gold_lists, alpha, verify_monotone=False) + except ValueError: + continue + + def admit(s, lam=lam): + return s <= lam + + n_ok_trials += 1 + hits, ngold = 0, 0 + uncal_hits, uncal_ngold = 0, 0 + eff_sum, raw_sum = 0.0, 0.0 + sentence_losses: List[float] = [] # CRC's own per-sentence loss + test_gold_source = combined_gold if pool_and_resplit else test_gold + test_probs_source = combined_probs if pool_and_resplit else test_pool.probs + test_cls_source = combined_id_to_class if pool_and_resplit else test_pool.id_to_class + for i in test_idx: + probs = test_probs_source[i] + cls_map = test_cls_source[i] + L, K, C = probs.shape + sentence_gold = [(t, s) for t, s in test_gold_source[i] if t in calibrated_types] + sentence_hits = 0 + for etype, s in test_gold_source[i]: + if etype in calibrated_types: + ngold += 1 + per_type_n[etype] = per_type_n.get(etype, 0) + 1 + if admit(s): + hits += 1 + sentence_hits += 1 + per_type_hits[etype] = per_type_hits.get(etype, 0) + 1 + else: + # descriptive only, no guarantee -- raw p>0.5 rule, matching + # ConformalGLiNER's own out-of-calibration fallback behavior. + uncal_ngold += 1 + if s <= 0.5: + uncal_hits += 1 + # CRC's own loss convention: 0 for entity-free sentences, avoids a 0/0 + # and matches exactly what crc_lambda_search calibrated against. + sentence_losses.append(1.0 - sentence_hits / len(sentence_gold) if sentence_gold else 0.0) + for col in range(C): + etype = cls_map.get(col + 1) + if etype is None or etype not in calibrated_types: + continue + nc = 1.0 - probs[:, :, col] + thresh = tau if mode == "span_filter" else lam + eff_sum += (nc <= thresh).sum().item() + raw_sum += L * K + + if mode == "risk_control": + # Report the quantity CRC actually calibrates and guarantees: the mean + # PER-SENTENCE miss rate, not entities pooled flat across sentences. + # These differ whenever gold-entity count per sentence is uneven -- + # pooling flat would silently measure a different, uncalibrated + # quantity and can show spurious "undercoverage" that has nothing to + # do with the (valid) CRC guarantee actually being tested. + coverages.append(1.0 - sum(sentence_losses) / len(sentence_losses) if sentence_losses else float("nan")) + else: + coverages.append(hits / ngold if ngold else float("nan")) + if uncal_ngold: + uncal_coverages.append(uncal_hits / uncal_ngold) + effs.append(eff_sum / len(test_idx)) + raw_counts.append(raw_sum / len(test_idx)) + + per_type_coverage = {t: per_type_hits.get(t, 0) / n for t, n in per_type_n.items() if n > 0} + return { + "alpha": alpha, + "n_calib": n_calib, + "n_trials_requested": n_trials, + "n_trials_ok": n_ok_trials, + "uncalibrated_coverage_mean": (sum(uncal_coverages) / len(uncal_coverages)) if uncal_coverages else None, + "n_uncalibrated_trials_with_data": len(uncal_coverages), + "coverage_mean": sum(coverages) / len(coverages) if coverages else float("nan"), + "coverage_std": (sum((c - sum(coverages) / len(coverages)) ** 2 for c in coverages) / len(coverages)) ** 0.5 + if coverages + else float("nan"), + "efficiency_mean": sum(effs) / len(effs) if effs else float("nan"), + "raw_candidates_mean": sum(raw_counts) / len(raw_counts) if raw_counts else float("nan"), + "per_type_coverage": per_type_coverage, + } + + +def run_suite( + name: str, calib_pool: Pool, test_pool: Pool, n_trials: int, n_calib: int, pool_and_resplit: bool = False +) -> List[Dict]: + calib_gold = gold_nc_by_example(calib_pool) + test_gold = gold_nc_by_example(test_pool) + rows = [] + for mode in ("span_filter", "risk_control"): + for alpha in ALPHAS: + t0 = time.time() + m = trial_metrics( + calib_pool, + test_pool, + calib_gold, + test_gold, + alpha, + n_calib, + n_trials, + mode, + seed=1234, + pool_and_resplit=pool_and_resplit, + ) + m.update({"pair": name, "mode": mode, "seconds": round(time.time() - t0, 1)}) + rows.append(m) + uncal = m["uncalibrated_coverage_mean"] + uncal_str = f", uncalibrated_coverage={uncal:.4f} (no guarantee)" if uncal is not None else "" + print( + f"[{name}/{mode}/alpha={alpha}] coverage={m['coverage_mean']:.4f}" + f"+-{m['coverage_std']:.4f} eff={m['efficiency_mean']:.1f}" + f" ({m['n_trials_ok']}/{n_trials} trials, {m['seconds']}s){uncal_str}" + ) + return rows + + +def calib_size_sensitivity(calib_pool: Pool, test_pool: Pool, alpha: float, n_trials: int) -> List[Dict]: + calib_gold = gold_nc_by_example(calib_pool) + test_gold = gold_nc_by_example(test_pool) + rows = [] + for n_calib in [50, 100, 200, 500, 1000]: + if n_calib > len(calib_pool.examples): + continue + m = trial_metrics( + calib_pool, + test_pool, + calib_gold, + test_gold, + alpha, + n_calib, + n_trials, + "span_filter", + 99, + pool_and_resplit=True, + ) + m["n_calib"] = n_calib + rows.append(m) + print(f"[calib_size n={n_calib}] coverage={m['coverage_mean']:.4f}+-{m['coverage_std']:.4f}") + return rows + + +def make_plots(rows: List[Dict], sensitivity_rows: List[Dict], out_dir: Path) -> None: + out_dir.mkdir(parents=True, exist_ok=True) + + # (a) coverage vs alpha, small multiples per (pair, mode) + pairs_modes = sorted({(r["pair"], r["mode"]) for r in rows}) + fig, axes = plt.subplots(1, len(pairs_modes), figsize=(5 * len(pairs_modes), 4), sharey=True) + if len(pairs_modes) == 1: + axes = [axes] + for ax, (pair, mode) in zip(axes, pairs_modes): + sub = sorted([r for r in rows if r["pair"] == pair and r["mode"] == mode], key=lambda r: r["alpha"]) + xs = [r["alpha"] for r in sub] + ys = [r["coverage_mean"] for r in sub] + es = [r["coverage_std"] for r in sub] + ax.errorbar(xs, ys, yerr=es, marker="o", label="empirical") + ax.plot([0, 1], [1, 0], "k--", alpha=0.5, label="y=1-alpha") + ax.set_xlim(0, 0.25) + ax.set_ylim(0, 1.05) + ax.set_title(f"{pair}\n{mode}") + ax.set_xlabel("alpha") + axes[0].set_ylabel("empirical coverage") + axes[0].legend() + fig.tight_layout() + fig.savefig(out_dir / "coverage_vs_alpha.png", dpi=150) + plt.close(fig) + + # (b) efficiency vs alpha + fig, axes = plt.subplots(1, len(pairs_modes), figsize=(5 * len(pairs_modes), 4), sharey=False) + if len(pairs_modes) == 1: + axes = [axes] + for ax, (pair, mode) in zip(axes, pairs_modes): + sub = sorted([r for r in rows if r["pair"] == pair and r["mode"] == mode], key=lambda r: r["alpha"]) + xs = [r["alpha"] for r in sub] + ys = [r["efficiency_mean"] for r in sub] + raw = [r["raw_candidates_mean"] for r in sub] + ax.plot(xs, ys, marker="o", label="admitted (efficiency)") + ax.plot(xs, raw, "k--", alpha=0.5, label="raw candidates (pre-filter)") + ax.set_title(f"{pair}\n{mode}") + ax.set_xlabel("alpha") + ax.set_yscale("log") + axes[0].set_ylabel("mean candidates / sentence") + axes[0].legend() + fig.tight_layout() + fig.savefig(out_dir / "efficiency_vs_alpha.png", dpi=150) + plt.close(fig) + + # (c) per-class coverage at alpha=0.1, span_filter mode + pairs = sorted({r["pair"] for r in rows}) + fig, axes = plt.subplots(1, len(pairs), figsize=(6 * len(pairs), 4)) + if len(pairs) == 1: + axes = [axes] + for ax, pair in zip(axes, pairs): + row = next((r for r in rows if r["pair"] == pair and r["mode"] == "span_filter" and r["alpha"] == 0.10), None) + if row is None: + continue + types = sorted(row["per_type_coverage"]) + vals = [row["per_type_coverage"][t] for t in types] + ax.bar(types, vals) + ax.axhline(0.9, color="k", linestyle="--", alpha=0.5) + ax.set_title(f"{pair} (alpha=0.1, span_filter)") + ax.set_ylim(0, 1.05) + ax.tick_params(axis="x", rotation=45) + fig.tight_layout() + fig.savefig(out_dir / "per_class_coverage.png", dpi=150) + plt.close(fig) + + # (d) calibration-size sensitivity + if sensitivity_rows: + fig, ax = plt.subplots(figsize=(6, 4)) + xs = [r["n_calib"] for r in sensitivity_rows] + ys = [r["coverage_mean"] for r in sensitivity_rows] + es = [r["coverage_std"] for r in sensitivity_rows] + ax.errorbar(xs, ys, yerr=es, marker="o") + ax.axhline(0.9, color="k", linestyle="--", alpha=0.5) + ax.set_xscale("log") + ax.set_xlabel("n_calib") + ax.set_ylabel("empirical coverage (alpha=0.1)") + ax.set_title("Calibration-set-size sensitivity (in-domain CoNLL-2003)") + fig.tight_layout() + fig.savefig(out_dir / "calib_size_sensitivity.png", dpi=150) + plt.close(fig) + + +def write_results_md(rows: List[Dict], sensitivity_rows: List[Dict], out_dir: Path, n_trials: int) -> None: + lines = [ + "# Conformal-GLiNER Empirical Validation Results", + "", + f"Model: `{MODEL_ID}`. Trials per (pair, mode, alpha): {n_trials}. Pool cap: {POOL_CAP} sentences.", + "", + "**Disclosed scope**: this run covers in-domain CoNLL-2003, in-domain WNUT-17, and a " + "zero-shot pair (calibrate on CoNLL-2003, measure coverage on WNUT-17). It does not " + "cover a full multi-domain CrossNER sweep -- documented as future work, not silently " + "dropped.", + "", + "## Summary table", + "", + "`coverage_mean` is over calibrated types only (the actually-guaranteed number). " + "`uncalibrated_coverage` (when present) is the raw p>0.5 empirical rate for types " + "requested at test time that never met the calibration floor -- descriptive only, " + "carries no guarantee, and is exactly what the zero-shot descope predicts will happen " + "for the WNUT-only types.", + "", + "| pair | mode | alpha | n_calib | trials_ok | coverage_mean | coverage_std " + "| efficiency_mean | raw_candidates_mean | uncalibrated_coverage |", + "|---|---|---|---|---|---|---|---|---|---|", + ] + for r in rows: + uncal = r["uncalibrated_coverage_mean"] + uncal_str = f"{uncal:.4f}" if uncal is not None else "n/a" + lines.append( + f"| {r['pair']} | {r['mode']} | {r['alpha']} | {r['n_calib']} " + f"| {r['n_trials_ok']}/{r['n_trials_requested']} " + f"| {r['coverage_mean']:.4f} | {r['coverage_std']:.4f} " + f"| {r['efficiency_mean']:.1f} | {r['raw_candidates_mean']:.1f} | {uncal_str} |" + ) + + lines += ["", "## Per-type coverage (span_filter, alpha=0.1, calibrated types only)", ""] + for r in rows: + if r["mode"] == "span_filter" and r["alpha"] == 0.10: + lines.append(f"**{r['pair']}**:") + for t, c in sorted(r["per_type_coverage"].items()): + lines.append(f"- {t}: {c:.4f}") + lines.append("") + + if sensitivity_rows: + lines += ["## Calibration-set-size sensitivity (in-domain CoNLL-2003, alpha=0.1, span_filter)", ""] + lines.append("| n_calib | coverage_mean | coverage_std |") + lines.append("|---|---|---|") + for r in sensitivity_rows: + lines.append(f"| {r['n_calib']} | {r['coverage_mean']:.4f} | {r['coverage_std']:.4f} |") + + (out_dir / "RESULTS.md").write_text("\n".join(lines)) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--output_dir", default="results/conformal") + ap.add_argument("--n_trials", type=int, default=50) + ap.add_argument("--n_calib", type=int, default=500) + ap.add_argument("--pool_cap", type=int, default=POOL_CAP) + args = ap.parse_args() + + out_dir = Path(args.output_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + print(f"Loading model {MODEL_ID}...") + model = GLiNER.from_pretrained(MODEL_ID) + + conll_labels = ["person", "organisation", "location", "misc"] + wnut_labels = ["corporation", "creative-work", "group", "location", "person", "product"] + + print("Loading CoNLL-2003 validation...") + conll_val = load_examples("DFKI-SLT/cross_ner", "conll2003", "validation", args.pool_cap) + print("Loading CoNLL-2003 test...") + conll_test = load_examples("DFKI-SLT/cross_ner", "conll2003", "test", args.pool_cap) + print("Loading WNUT-17 validation...") + wnut_val = load_examples("leondz/wnut_17", None, "validation", args.pool_cap) + print("Loading WNUT-17 test...") + wnut_test = load_examples("leondz/wnut_17", None, "test", args.pool_cap) + + print(f"CoNLL val/test: {len(conll_val)}/{len(conll_test)}, WNUT val/test: {len(wnut_val)}/{len(wnut_test)}") + + print("Building pools (forward passes)...") + t0 = time.time() + conll_calib_pool = build_pool(model, "conll_calib", conll_val, conll_labels) + conll_test_pool = build_pool(model, "conll_test", conll_test, conll_labels) + wnut_calib_pool = build_pool(model, "wnut_calib", wnut_val, wnut_labels) + wnut_test_pool = build_pool(model, "wnut_test", wnut_test, wnut_labels) + # Pair A: calibrate on CoNLL types, test on WNUT types -- score the WNUT test pool + # against WNUT's own label set (already have wnut_test_pool for that), and score the + # CoNLL calibration pool against CoNLL's own labels (already have conll_calib_pool). + print(f"Pools built in {time.time() - t0:.1f}s") + + rows = [] + rows += run_suite( + "in-domain CoNLL-2003", + conll_calib_pool, + conll_test_pool, + args.n_trials, + args.n_calib, + pool_and_resplit=True, + ) + rows += run_suite( + "in-domain WNUT-17", wnut_calib_pool, wnut_test_pool, args.n_trials, args.n_calib, pool_and_resplit=True + ) + rows += run_suite( + "zero-shot CoNLL-2003->WNUT-17 (Pair A)", conll_calib_pool, wnut_test_pool, args.n_trials, args.n_calib + ) + + print("Calibration-size sensitivity (in-domain CoNLL-2003)...") + sensitivity_rows = calib_size_sensitivity(conll_calib_pool, conll_test_pool, alpha=0.10, n_trials=args.n_trials) + + (out_dir / "raw_results.json").write_text(json.dumps({"rows": rows, "sensitivity": sensitivity_rows}, indent=2)) + make_plots(rows, sensitivity_rows, out_dir) + write_results_md(rows, sensitivity_rows, out_dir, args.n_trials) + print(f"Done. Results in {out_dir}/") + + +if __name__ == "__main__": + main() diff --git a/tests/test_conformal_calibrators.py b/tests/test_conformal_calibrators.py new file mode 100644 index 00000000..e5770fb2 --- /dev/null +++ b/tests/test_conformal_calibrators.py @@ -0,0 +1,142 @@ +"""Synthetic, network-free tests for gliner/conformal/calibrators.py. + +Mirrors tests/test_decoder.py's pattern: hand-built inputs with analytically +known ground truth, no model download. See docs/conformal.md for the theory +these tests check against. +""" + +import random + +import pytest + +from gliner.conformal.calibrators import ( + calibration_floor, + crc_lambda_search, + mondrian_calibrate, + split_conformal_quantile, +) + + +class TestCalibrationFloor: + def test_matches_eval_plan_table(self): + # Worked table: n >= ceil((1-alpha)/alpha). + assert calibration_floor(0.20) == 4 + assert calibration_floor(0.10) == 9 + assert calibration_floor(0.05) == 19 + + def test_rejects_invalid_alpha(self): + with pytest.raises(ValueError): + calibration_floor(0.0) + with pytest.raises(ValueError): + calibration_floor(1.0) + with pytest.raises(ValueError): + calibration_floor(-0.1) + + +class TestSplitConformalQuantile: + def test_raises_below_floor(self): + with pytest.raises(ValueError, match="insufficient"): + split_conformal_quantile([0.1, 0.2, 0.3, 0.4, 0.5], alpha=0.05) + + def test_at_exact_floor_returns_the_max(self): + floor = calibration_floor(0.2) + scores = [i / 10 for i in range(floor)] + assert split_conformal_quantile(scores, alpha=0.2) == max(scores) + + def test_empirical_coverage_matches_theory(self): + """20000 seeded trials: split-conformal coverage on Uniform(0,1) scores + should land within a few standard errors of the 1-alpha target.""" + rng = random.Random(42) + n, alpha, trials = 500, 0.1, 20000 + hits = 0 + for _ in range(trials): + calib = [rng.random() for _ in range(n)] + test = rng.random() + q = split_conformal_quantile(calib, alpha) + hits += test <= q + coverage = hits / trials + se = (coverage * (1 - coverage) / trials) ** 0.5 + target = 1 - alpha + assert target - 4 * se <= coverage <= target + 1 / (n + 1) + 4 * se + + def test_rejects_invalid_alpha(self): + with pytest.raises(ValueError): + split_conformal_quantile([0.1, 0.2], alpha=1.5) + + +class TestMondrianCalibrate: + def test_skips_sub_floor_types_and_calibrates_the_rest(self): + rng = random.Random(0) + alpha = 0.1 + floor = calibration_floor(alpha) + scores_by_type = { + "common": [rng.random() for _ in range(200)], + "rare": [rng.random() for _ in range(floor - 1)], + } + thresholds, skipped = mondrian_calibrate(scores_by_type, alpha) + assert "common" in thresholds + assert "rare" not in thresholds + assert skipped == {"rare": floor - 1} + + def test_empty_input(self): + thresholds, skipped = mondrian_calibrate({}, 0.1) + assert thresholds == {} + assert skipped == {} + + +class TestCrcLambdaSearch: + def test_raises_below_floor(self): + with pytest.raises(ValueError, match="insufficient"): + crc_lambda_search([[0.1], [0.2], [0.3]], alpha=0.05) + + def test_boundary_case_all_scores_zero_gives_lambda_zero(self): + # Every gold entity perfectly scored (nonconformity 0) -> even the + # tightest threshold (lambda=0) already achieves zero risk. + gold = [[0.0, 0.0] for _ in range(50)] + lam = crc_lambda_search(gold, alpha=0.1) + assert lam == 0.0 + + def test_unrepresentable_entities_do_not_block_convergence_when_rare(self): + # A structurally-unrepresentable gold entity (float("inf")) can never + # be covered. If only a small fraction of examples have one (each + # contributing a fixed loss-1 floor), the target is still reachable + # as long as that floor alone is below alpha. + rng = random.Random(1) + easy = [[rng.random() * 0.05] for _ in range(190)] + unrepresentable = [[float("inf")] for _ in range(10)] + gold = easy + unrepresentable + lam = crc_lambda_search(gold, alpha=0.2) + assert lam < float("inf") + + def test_unrepresentable_entities_correctly_block_convergence_when_common(self): + # If unrepresentable entities are common enough that even lambda=inf + # cannot bring the risk under alpha, returning inf (not a finite but + # invalid lambda) is the mathematically correct answer, not a bug. + rng = random.Random(1) + gold = [[rng.random() * 0.05, float("inf")] for _ in range(200)] + lam = crc_lambda_search(gold, alpha=0.2) + assert lam == float("inf") + + def test_empirical_risk_control_matches_theory(self): + """CRC's proved guarantee: E[miss_rate] <= alpha on fresh test data.""" + rng = random.Random(7) + alpha = 0.1 + gold_calib = [[rng.random()] for _ in range(1000)] + lam = crc_lambda_search(gold_calib, alpha) + + trials, n_test = 200, 500 + miss_rates = [] + for _ in range(trials): + test = [[rng.random()] for _ in range(n_test)] + missed = sum(1 for g in test if g[0] > lam) / n_test + miss_rates.append(missed) + mean_miss = sum(miss_rates) / len(miss_rates) + assert mean_miss <= alpha + 0.02 # small slack for Monte Carlo noise + + def test_monotonicity_precondition_is_checked_by_default(self): + # Sanity: verify_monotone=True must not raise on a genuinely nested + # (by construction) family -- this is the runtime check on CRC's + # monotonicity precondition. + rng = random.Random(3) + gold = [[rng.random() for _ in range(rng.randint(0, 3))] for _ in range(100)] + crc_lambda_search(gold, alpha=0.2, verify_monotone=True) # must not raise diff --git a/tests/test_conformal_gliner.py b/tests/test_conformal_gliner.py new file mode 100644 index 00000000..66578885 --- /dev/null +++ b/tests/test_conformal_gliner.py @@ -0,0 +1,349 @@ +"""Integration tests for ConformalGLiNER against a real small checkpoint. + +Mirrors tests/test_models.py::test_span_model's pattern (the only other test +in the suite that downloads a real model, gliner-community/gliner_small-v2.5). +This is the only conformal test module that touches the network; +tests/test_conformal_calibrators.py is fully synthetic. +""" + +import warnings + +import pytest + +from gliner import GLiNER +from gliner.conformal import ConformalGLiNER, align_gold_scores, extract_raw_scores +from gliner.conformal.calibrators import calibration_floor +from gliner.conformal.scores import _assert_span_mode_supported + +MODEL_ID = "gliner-community/gliner_small-v2.5" + + +def _examples(n_per_type: int = 10): + """Small, easy, synthetic calibration/test corpus with a fixed 3-type schema.""" + templates = [ + ( + "Apple was founded by Steve Jobs in Cupertino .", + [(0, 0, "organization"), (4, 5, "person"), (7, 7, "location")], + ), + ( + "Google was founded by Larry Page in California .", + [(0, 0, "organization"), (4, 5, "person"), (7, 7, "location")], + ), + ( + "Microsoft was founded by Bill Gates in Redmond .", + [(0, 0, "organization"), (4, 5, "person"), (7, 7, "location")], + ), + ( + "Amazon was founded by Jeff Bezos in Seattle .", + [(0, 0, "organization"), (4, 5, "person"), (7, 7, "location")], + ), + ("Tesla was founded by Elon Musk in Austin .", [(0, 0, "organization"), (4, 5, "person"), (7, 7, "location")]), + ( + "IBM was founded by Charles Flint in New York .", + [(0, 0, "organization"), (4, 5, "person"), (7, 8, "location")], + ), + ( + "Intel was founded by Robert Noyce in Santa Clara .", + [(0, 0, "organization"), (4, 5, "person"), (7, 8, "location")], + ), + ( + "Oracle was founded by Larry Ellison in Redwood City .", + [(0, 0, "organization"), (4, 5, "person"), (7, 8, "location")], + ), + ] + out = [] + i = 0 + while len(out) < n_per_type: + text, ner = templates[i % len(templates)] + out.append({"tokenized_text": text.split(), "ner": [list(t) for t in ner]}) + i += 1 + return out + + +@pytest.fixture(scope="module") +def model(): + return GLiNER.from_pretrained(MODEL_ID) + + +@pytest.fixture(scope="module") +def calib_data(): + # 25 examples per type comfortably clears calibration_floor(0.2) == 4 and + # calibration_floor(0.1) == 9, used throughout this module. + return _examples(25) + + +class TestCalibrateAllModes: + @pytest.mark.parametrize("mode", ["span_filter", "risk_control", "mondrian"]) + def test_calibrate_and_predict_smoke(self, model, calib_data, mode): + cg = ConformalGLiNER(model) + cg.calibrate(calib_data, alpha=0.2, mode=mode) + assert cg.is_calibrated + assert set(cg._state.calibrated_types) == {"organization", "person", "location"} + + preds = cg.predict_entities( + "Netflix was founded by Reed Hastings in Los Gatos .", ["organization", "person", "location"] + ) + assert isinstance(preds, list) + for ent in preds: + assert ent["conformal"]["mode"] == mode + assert ent["conformal"]["calibrated"] is True + + def test_batch_predict_shape(self, model, calib_data): + cg = ConformalGLiNER(model).calibrate(calib_data, alpha=0.2, mode="risk_control") + preds = cg.predict_entities( + ["Netflix was founded by Reed Hastings .", "Uber was founded by Travis Kalanick ."], + ["organization", "person"], + ) + assert isinstance(preds, list) and len(preds) == 2 + assert all(isinstance(p, list) for p in preds) + + +class TestCalibrationFloorEnforcement: + def test_raises_with_too_few_examples(self, model): + cg = ConformalGLiNER(model) + tiny = _examples(2) # below calibration_floor(0.05) == 19 + with pytest.raises(ValueError, match=r"floor|insufficient"): + cg.calibrate(tiny, alpha=0.05, mode="risk_control") + + def test_warns_for_under_floor_type_but_still_calibrates_others(self, model): + floor = calibration_floor(0.1) + assert floor == 9 + data = _examples(floor + 5) # organization/person/location all clear the floor + # Add a handful of a fourth type that stays under floor. + data.append({"tokenized_text": ["Rare", "Corp", "makes", "widgets", "."], "ner": [[0, 1, "rare_type"]]}) + cg = ConformalGLiNER(model) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + cg.calibrate(data, alpha=0.1, mode="span_filter") + assert any("rare_type" in str(w.message) for w in caught) + assert "rare_type" not in cg._state.calibrated_types + assert {"organization", "person", "location"} <= set(cg._state.calibrated_types) + + +class TestUncalibratedTypeFallback: + def test_unseen_label_warns_and_is_flagged_uncalibrated(self, model, calib_data): + cg = ConformalGLiNER(model).calibrate(calib_data, alpha=0.2, mode="risk_control") + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + preds = cg.predict_entities( + "The chemical compound was synthesized in the lab .", + ["organization", "chemical_compound"], + ) + assert any("chemical_compound" in str(w.message) for w in caught) + for ent in preds: + if ent["label"] == "chemical_compound": + assert ent["conformal"]["calibrated"] is False + + +class TestEmptyPredictions: + def test_no_matching_entities_returns_empty_list(self, model, calib_data): + cg = ConformalGLiNER(model).calibrate(calib_data, alpha=0.2, mode="risk_control") + preds = cg.predict_entities("zzz qqq xxx yyy .", ["organization", "person", "location"]) + assert preds == [] + + def test_empty_text_returns_empty(self, model, calib_data): + cg = ConformalGLiNER(model).calibrate(calib_data, alpha=0.2, mode="risk_control") + preds = cg.predict_entities("", ["organization"]) + assert preds == [] + + +class TestSaveLoadRoundTrip: + def test_round_trip(self, model, calib_data, tmp_path): + cg = ConformalGLiNER(model).calibrate(calib_data, alpha=0.2, mode="risk_control") + path = tmp_path / "calibration.json" + cg.save_calibration(str(path)) + + cg2 = ConformalGLiNER.load_calibration(str(path), model) + assert cg2._state.mode == cg._state.mode + assert cg2._state.alpha == cg._state.alpha + assert cg2._state.crc_lambda == cg._state.crc_lambda + + text = "Netflix was founded by Reed Hastings in Los Gatos ." + labels = ["organization", "person", "location"] + assert cg.predict_entities(text, labels) == cg2.predict_entities(text, labels) + + def test_load_warns_on_model_mismatch(self, model, calib_data, tmp_path): + cg = ConformalGLiNER(model).calibrate(calib_data, alpha=0.2, mode="risk_control") + path = tmp_path / "calibration.json" + cg.save_calibration(str(path)) + cg._state.model_id = "some/other-model" + cg.save_calibration(str(path)) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + ConformalGLiNER.load_calibration(str(path), model) + assert any("model-specific" in str(w.message) or "Nonconformity" in str(w.message) for w in caught) + + +class TestCoverageReport: + def test_report_shape_and_disjoint_data_canary(self, model, calib_data): + cg = ConformalGLiNER(model).calibrate(calib_data, alpha=0.2, mode="risk_control") + + # Canary: coverage measured on the *same* data the + # threshold was calibrated on must come out at or above the nominal + # target, since the threshold was tuned to fit exactly this data -- + # a biased estimate, and this test documents/guards that property + # rather than treating it as a valid held-out coverage number. + report = cg.coverage_report(calib_data) + assert report["overall_coverage"] >= 1 - cg._state.alpha - 1e-9 + assert report["n_uncalibrated_gold"] == 0 + assert set(report["per_type_coverage"]) == {"organization", "person", "location"} + assert report["efficiency_mean"] >= 0 + assert report["raw_candidates_mean"] > 0 + + def test_risk_control_reports_per_sentence_not_per_entity_pooled(self, model, calib_data): + """Regression test for a real bug found during empirical validation: + risk_control calibrates and guarantees a *per-sentence* average miss + rate (Conformal Risk Control's own loss definition), which is a + different quantity from pooling every gold entity flat across + sentences whenever entity-count-per-sentence varies. A test corpus + with 1 entity in one sentence and 3 in another + makes the two quantities provably different, so a regression back to + flat pooling shows up as a hard assertion failure, not a subtle + drift in a coverage number.""" + cg = ConformalGLiNER(model).calibrate(calib_data, alpha=0.2, mode="risk_control") + + sentence_one = ["Apple", "was", "founded", "by", "Steve", "Jobs", "in", "Cupertino", "."] + sentence_two = ["Google", "Microsoft", "Amazon", "dominate", "the", "market", "."] + test_data = [ + {"tokenized_text": sentence_one, "ner": [[0, 0, "organization"]]}, + { + "tokenized_text": sentence_two, + "ner": [[0, 0, "organization"], [1, 1, "organization"], [2, 2, "organization"]], + }, + ] + report = cg.coverage_report(test_data) + + raw = extract_raw_scores(model, test_data, ["organization"]) + scores, types, example_idx = align_gold_scores(raw, test_data) + tau = cg._nc_threshold_for(cg._state, "organization") + hits_by_example = {0: [], 1: []} + for s, _t, i in zip(scores, types, example_idx): + hits_by_example[i].append(s <= tau) + + pooled = sum(sum(h) for h in hits_by_example.values()) / sum(len(h) for h in hits_by_example.values()) + per_sentence = sum((sum(h) / len(h) if h else 1.0) for h in hits_by_example.values()) / len(hits_by_example) + + assert report["overall_coverage"] == pytest.approx(per_sentence) + if pooled != per_sentence: + assert report["overall_coverage"] != pytest.approx(pooled) + + +class TestRequiresCalibration: + def test_predict_before_calibrate_raises(self, model): + cg = ConformalGLiNER(model) + with pytest.raises(RuntimeError, match="not calibrated"): + cg.predict_entities("Apple was founded by Steve Jobs .", ["organization"]) + + def test_coverage_report_before_calibrate_raises(self, model, calib_data): + cg = ConformalGLiNER(model) + with pytest.raises(RuntimeError, match="not calibrated"): + cg.coverage_report(calib_data) + + def test_thresholds_before_calibrate_raises(self, model): + cg = ConformalGLiNER(model) + with pytest.raises(RuntimeError, match="not calibrated"): + cg.thresholds() + + def test_calibrated_types_before_calibrate_raises(self, model): + cg = ConformalGLiNER(model) + with pytest.raises(RuntimeError, match="not calibrated"): + _ = cg.calibrated_types + + +class TestPublicThresholdAPI: + """Regression coverage for the per-label threshold API requested in PR review + (urchade/GLiNER#374) -- exposing what was previously only reachable via the + private ``_state`` attribute.""" + + @pytest.mark.parametrize("mode", ["span_filter", "risk_control", "mondrian"]) + def test_calibrated_types_matches_internal_state(self, model, calib_data, mode): + cg = ConformalGLiNER(model).calibrate(calib_data, alpha=0.2, mode=mode) + assert set(cg.calibrated_types) == {"organization", "person", "location"} + # public accessor, not a live reference to internal state + cg.calibrated_types.append("tampered") + assert "tampered" not in cg.calibrated_types + + @pytest.mark.parametrize("mode", ["span_filter", "risk_control", "mondrian"]) + def test_thresholds_covers_every_calibrated_type(self, model, calib_data, mode): + cg = ConformalGLiNER(model).calibrate(calib_data, alpha=0.2, mode=mode) + thresholds = cg.thresholds() + assert set(thresholds.keys()) == set(cg.calibrated_types) + assert all(isinstance(v, float) for v in thresholds.values()) + + def test_mondrian_thresholds_can_differ_per_type(self, model, calib_data): + cg = ConformalGLiNER(model).calibrate(calib_data, alpha=0.2, mode="mondrian") + thresholds = cg.thresholds() + # Not asserting they DO differ (real calibration data may coincidentally + # produce equal thresholds) -- asserting the API *can* express a + # per-type difference, unlike span_filter/risk_control below. + assert isinstance(thresholds, dict) and len(thresholds) == 3 + + @pytest.mark.parametrize("mode", ["span_filter", "risk_control"]) + def test_pooled_modes_share_one_threshold_across_labels(self, model, calib_data, mode): + cg = ConformalGLiNER(model).calibrate(calib_data, alpha=0.2, mode=mode) + thresholds = cg.thresholds() + assert len(set(thresholds.values())) == 1 + + +class TestModelCalibrateConvenienceMethod: + """Regression coverage for the `model.calibrate()` / `model.conformal` API + requested in PR review (urchade/GLiNER#374) -- calibration and inference on + the same object, not just through a separately-constructed ConformalGLiNER. + + ``model`` is a module-scoped fixture shared across this whole test file -- + every test here must undo its own calibration afterward so it doesn't leak + into unrelated tests that assume an uncalibrated model.""" + + @pytest.fixture + def calibrated_model(self, model, calib_data): + model.calibrate(calib_data, alpha=0.2, mode="risk_control") + yield model + model._conformal_model = None + + def test_conformal_is_none_before_calibrate(self, model): + assert model.conformal is None + + def test_calibrate_returns_self_for_chaining(self, model, calib_data): + try: + result = model.calibrate(calib_data, alpha=0.2, mode="risk_control") + assert result is model + finally: + model._conformal_model = None + + def test_conformal_property_exposes_a_calibrated_wrapper(self, calibrated_model): + assert isinstance(calibrated_model.conformal, ConformalGLiNER) + assert calibrated_model.conformal.is_calibrated + assert set(calibrated_model.conformal.calibrated_types) == {"organization", "person", "location"} + + def test_predicting_through_the_stored_wrapper_matches_direct_wrapper_use(self, calibrated_model): + text = "Netflix was founded by Reed Hastings in Los Gatos ." + labels = ["organization", "person", "location"] + + via_model = calibrated_model.conformal.predict_entities(text, labels) + + fresh_wrapper = ConformalGLiNER(calibrated_model) + fresh_wrapper._state = calibrated_model.conformal._state # same calibration, no re-fitting + via_fresh_wrapper = fresh_wrapper.predict_entities(text, labels) + + assert via_model == via_fresh_wrapper + + def test_unsupported_architecture_raises_not_implemented(self): + class _FakeTokenModel: + pass + + # calibrate() is only meaningful on real BaseEncoderGLiNER instances; + # this documents that the NotImplementedError comes from ConformalGLiNER + # itself (see TestTokenModeRejected), not duplicated validation here. + with pytest.raises(NotImplementedError, match="span-mode"): + _assert_span_mode_supported(_FakeTokenModel()) + + +class TestTokenModeRejected: + def test_non_span_mode_model_raises_not_implemented(self): + class _FakeTokenModel: + pass + + cg = ConformalGLiNER(_FakeTokenModel()) + with pytest.raises(NotImplementedError, match="span-mode"): + cg.calibrate(_examples(20), alpha=0.2, mode="risk_control")