hallustream 0.2.0__py3-none-any.whl

hallustream/__init__.py

@@ -0,0 +1,18 @@
+"""HalluStream: flag the facts an open LLM is probably making up, from its hidden states."""
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from hallustream.detector import Detector, Entity, Event, Report
+
+__version__ = "0.2.0"
+__all__ = ["Detector", "Entity", "Event", "Report", "__version__"]
+
+
+def __getattr__(name: str):
+    # Loaded on first use, so the research code can import hallustream.* without the library
+    if name in ("Detector", "Entity", "Event", "Report"):
+        from hallustream import detector
+
+        return getattr(detector, name)
+    raise AttributeError(name)

hallustream/cli.py

@@ -0,0 +1,220 @@
+"""The hallustream command: answers with the details the model is probably making up highlighted.
+
+    hallustream "Tell me about Lake Tahoe."               ask one question (the default)
+    hallustream chat                                      have a conversation
+    hallustream score "Tell me about Lake Tahoe." a.txt   check an answer it already wrote
+    hallustream batch prompts.jsonl -o results.jsonl      check many prompts or answers
+
+Answers appear sentence by sentence, once each sentence has been checked.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+from hallustream import __version__
+from hallustream.detector import BUNDLED, Detector, Entity, Report, _weights_dir, highlight
+
+COMMANDS = ("ask", "chat", "score", "batch")
+DIM, RESET = "\033[2m", "\033[0m"
+
+
+def _parser() -> argparse.ArgumentParser:
+    shared = argparse.ArgumentParser(add_help=False)
+    shared.add_argument("--threshold", default="balanced",
+                        help="balanced (default), strict, or a risk between 0 and 1")  # fmt: skip
+    shared.add_argument("--weights", default="qwen3-4b", help="bundled name or weights folder")
+    shared.add_argument("--device", choices=["cuda", "mps", "cpu"],
+                        help="default: the fastest available")  # fmt: skip
+    p = argparse.ArgumentParser(
+        prog="hallustream",
+        description="Answers from Qwen3-4B with the details it is probably making up highlighted."
+        " Without a command, the arguments go to 'ask'.",
+        epilog='Example: hallustream "Tell me about Lake Tahoe."  (see "hallustream ask -h" for'
+        " --threshold, --device, --json and the other options)",
+    )
+    p.add_argument("--version", action="version", version=f"hallustream {__version__}")
+    sub = p.add_subparsers(dest="command", required=True)
+    ask = sub.add_parser("ask", parents=[shared], help="answer one question (the default)")
+    ask.add_argument("prompt")
+    chat = sub.add_parser("chat", parents=[shared], help="have a conversation")
+    score = sub.add_parser("score", parents=[shared], help="check an answer it already wrote")
+    score.add_argument("prompt", help="the question the answer was written for")
+    score.add_argument("file", nargs="?", default="-", help="the answer (default: stdin)")
+    batch = sub.add_parser("batch", parents=[shared], help="check a file of prompts or answers")
+    batch.add_argument("input", help='JSONL with "prompt" (and "answer" to check one already '
+                       'written), or one prompt per line; - for stdin')  # fmt: skip
+    batch.add_argument("-o", "--output", default="-", help="JSONL results (default: stdout)")
+    for s in (ask, chat, batch):
+        s.add_argument("--max-new-tokens", type=int, default=512)
+    for s in (ask, chat, score):
+        s.add_argument("--no-color", action="store_true")
+    for s in (ask, score):
+        s.add_argument("--json", action="store_true", help="print one JSON object instead")
+    return p
+
+
+def _load(weights: str, device: str | None) -> Detector:
+    import logging as pylogging
+    import warnings
+
+    import torch
+    from huggingface_hub import try_to_load_from_cache
+    from transformers.utils import logging
+
+    if device == "cuda" and not torch.cuda.is_available():
+        sys.exit("hallustream: --device cuda, but no CUDA GPU is available")
+    if device == "mps" and not torch.backends.mps.is_available():
+        sys.exit("hallustream: --device mps, but this machine has no Apple GPU")
+    # This is a command for reading answers: library warnings (GLiNER, transformers, torch's
+    # macOS note about redirects) are turned off for the whole process.
+    warnings.simplefilter("ignore")
+    pylogging.getLogger("torch.distributed.elastic.multiprocessing.redirects").setLevel("ERROR")
+
+    hf_id = json.loads((_weights_dir(weights) / "model.json").read_text())["hf_id"]
+    if isinstance(try_to_load_from_cache(hf_id, "model.safetensors.index.json"), str):
+        logging.disable_progress_bar()  # already downloaded: keep the terminal for the answer
+    else:
+        print(f"Downloading {hf_id} (about 8 GB, once)...", file=sys.stderr)
+    logging.set_verbosity_error()
+    return Detector.from_pretrained(weights, device=device)
+
+
+def _segment(text: str, lo: int, hi: int, flagged: list[Entity], color: bool) -> str:
+    inside = [Entity(e.text, e.start - lo, e.end - lo, e.risk, True)
+              for e in flagged if lo <= e.start and e.end <= hi]  # fmt: skip
+    return highlight(text[lo:hi], inside, "ansi" if color else "plain")
+
+
+def _stream(det: Detector, messages: list[dict], max_new_tokens: int, threshold, color: bool,
+            ) -> Report:  # fmt: skip
+    """Print each sentence once it has been checked; return the whole answer's report."""
+    text, printed, found, n = "", 0, [], 0
+    for ev in det.stream(messages=messages, max_new_tokens=max_new_tokens, threshold=threshold):
+        text += ev.text
+        found += ev.entities
+        n = len(ev.ids)
+        if ev.checked > printed:
+            print(_segment(text, printed, ev.checked, [e for e in found if e.flagged], color),
+                  end="", flush=True)  # fmt: skip
+            printed = ev.checked
+    cut = f"\n(cut off at {max_new_tokens} tokens; raise --max-new-tokens for more)"
+    print(_dim(cut, color) if n >= max_new_tokens else "")
+    return Report(text, found, det._threshold(threshold))
+
+
+def _dim(text: str, color: bool) -> str:
+    return f"{DIM}{text}{RESET}" if color else text
+
+
+def _legend(threshold: float, color: bool) -> str:
+    return _dim(f"Risk is the estimated chance a detail is made up; flagged from {threshold:.0%}.",
+                color)  # fmt: skip
+
+
+def _ask(det: Detector, a, threshold, color: bool) -> int:
+    if a.json:
+        report = det.generate(a.prompt, max_new_tokens=a.max_new_tokens, threshold=threshold)
+        print(json.dumps({"prompt": a.prompt, **report.to_dict()}, ensure_ascii=False))
+        return 0
+    msgs = [{"role": "user", "content": a.prompt}]
+    report = _stream(det, msgs, a.max_new_tokens, threshold, color)
+    print(f"\n{report.summary()}\n{_legend(report.threshold, color)}")
+    return 0
+
+
+def _chat(det: Detector, a, threshold, color: bool) -> int:
+    legend = _legend(det._threshold(threshold), color)
+    print(f"Chat with the model; likely made-up details are highlighted.\n{legend}\n"
+          "Type exit or press Ctrl-D to quit.")  # fmt: skip
+    messages: list[dict] = []
+    while True:
+        try:
+            line = input("\n> ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print()
+            return 0
+        if line.lower() in ("exit", "quit"):
+            return 0
+        if not line:
+            continue
+        messages.append({"role": "user", "content": line})
+        report = _stream(det, messages, a.max_new_tokens, threshold, color)
+        messages.append({"role": "assistant", "content": report.text})
+        print(f"\n{report.summary()}")
+
+
+def _score(det: Detector, a, threshold, color: bool) -> int:
+    text = (sys.stdin.read() if a.file == "-" else Path(a.file).read_text()).rstrip()
+    report = det.score(a.prompt, text, threshold=threshold)
+    if a.json:
+        print(json.dumps({"prompt": a.prompt, **report.to_dict()}, ensure_ascii=False))
+        return 0
+    print(report.highlight("ansi" if color else "plain"))
+    print(f"\n{report.summary()}\n{_legend(report.threshold, color)}")
+    return 0
+
+
+def _batch(det: Detector, a, threshold, color: bool) -> int:
+    raw = sys.stdin.read() if a.input == "-" else Path(a.input).read_text()
+    lines = [x for x in raw.splitlines() if x.strip()]
+    out = sys.stdout if a.output == "-" else open(a.output, "w")
+    for i, line in enumerate(lines, 1):
+        try:
+            item: dict = (json.loads(line) if line.lstrip().startswith("{")
+                          else {"prompt": line.strip()})  # fmt: skip
+        except json.JSONDecodeError as e:
+            sys.exit(f"{a.input} line {i}: not valid JSON ({e.msg})")
+        prompt, msgs = item.get("prompt"), item.get("messages")
+        if prompt is None and msgs is None:
+            sys.exit(f'{a.input} line {i}: needs a "prompt" or "messages"')
+        if "answer" in item:
+            report = det.score(prompt, item["answer"], messages=msgs, threshold=threshold)
+        else:
+            report = det.generate(prompt, messages=msgs, max_new_tokens=a.max_new_tokens,
+                                  threshold=threshold)  # fmt: skip
+        out.write(json.dumps({**item, **report.to_dict()}, ensure_ascii=False) + "\n")
+        out.flush()
+        print(f"{i}/{len(lines)}: {len(report.flagged)} flagged", file=sys.stderr)
+    if out is not sys.stdout:
+        out.close()
+    return 0
+
+
+def main(argv: list[str] | None = None, detector: Detector | None = None) -> int:
+    argv = list(sys.argv[1:] if argv is None else argv)
+    if argv and argv[0] not in (*COMMANDS, "-h", "--help", "--version"):
+        argv.insert(0, "ask")
+    p = _parser()
+    a = p.parse_args(argv)
+    if detector is None and not (_weights_dir(a.weights) / "model.json").exists():
+        p.error(f"--weights {a.weights}: not a bundled name {list(BUNDLED)} or a weights folder")
+    try:
+        threshold: str | float = float(a.threshold)
+        if not 0 <= threshold <= 1:
+            p.error(f"--threshold {a.threshold}: a risk must be between 0 and 1")
+    except ValueError:
+        threshold = a.threshold
+        file = _weights_dir(a.weights) / "thresholds.json"
+        presets = detector.thresholds if detector else json.loads(file.read_text())
+        if not isinstance(presets.get(threshold), dict):
+            names = [k for k, v in presets.items() if isinstance(v, dict)]
+            p.error(f"unknown threshold {threshold!r}; use one of {names} or a number")
+    path = getattr(a, "file", None) or getattr(a, "input", None)
+    if path and path != "-" and not Path(path).is_file():
+        p.error(f"{path}: no such file")
+    color = not getattr(a, "no_color", True) and sys.stdout.isatty()
+    run = {"ask": _ask, "chat": _chat, "score": _score, "batch": _batch}[a.command]
+    try:
+        det = detector or _load(a.weights, a.device)
+        return run(det, a, threshold, color)
+    except KeyboardInterrupt:
+        print(file=sys.stderr)
+        return 130
+
+
+if __name__ == "__main__":
+    sys.exit(main())

hallustream/detector.py

@@ -0,0 +1,315 @@
+"""Flag the facts a model is probably making up, after it writes or while it writes.
+
+    det = Detector.from_pretrained("qwen3-4b")
+    report = det.generate("Tell me about Lake Tahoe.")        # write an answer and check it
+    report = det.score("Tell me about Lake Tahoe.", answer)   # or check one it already wrote
+    print(report.highlight("ansi"), report.summary(), sep="\n\n")
+    for event in det.stream("Tell me about Lake Tahoe."):    # flags as each sentence ends
+        ...
+
+The probe only knows what its own model knows: score text written by the model it was
+trained for (Qwen3-4B-Instruct-2507 for the bundled weights), not other models' text.
+"""
+
+from __future__ import annotations
+
+import html
+import json
+import math
+from collections.abc import Callable, Iterator
+from dataclasses import asdict, dataclass, field
+from functools import partial
+from importlib.resources import files
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from hallustream import entities as _entities
+from hallustream.probe import ADAPTER, Forward, Probe, char_to_tokens, forward
+from hallustream.text import PRONOUNS, is_refusal, sentences
+
+BUNDLED = ("qwen3-4b",)
+SENTENCE_END = ".!?"
+EntityFinder = Callable[[str, list[tuple[int, int]]], list[tuple[int, int]]]
+
+
+@dataclass
+class Entity:
+    text: str
+    start: int  # character offsets in the answer
+    end: int
+    risk: float  # probability that the entity is made up (sigmoid of the adapted probe's score)
+    flagged: bool
+
+
+@dataclass
+class Report:
+    text: str
+    entities: list[Entity]
+    threshold: float
+
+    @property
+    def flagged(self) -> list[Entity]:
+        return [e for e in self.entities if e.flagged]
+
+    def highlight(self, style: str = "plain") -> str:
+        """The text with flagged entities marked: "plain" [[...]], "ansi" colors, or "html"."""
+        return highlight(self.text, self.flagged, style)
+
+    def summary(self) -> str:
+        """The flagged details, riskiest first, each listed once."""
+        checked = {e.text for e in self.entities}
+        risky: dict[str, float] = {}
+        for e in self.flagged:
+            risky[e.text] = max(risky.get(e.text, 0.0), e.risk)
+        if not checked:
+            return "No specific details to check."
+        if not risky:
+            return f"No details flagged ({len(checked)} checked)."
+        lines = [f"{len(risky)} of {len(checked)} details are likely made up:"]
+        lines += [f"  {r:4.0%}  {t}" for t, r in sorted(risky.items(), key=lambda x: -x[1])]
+        return "\n".join(lines)
+
+    def to_dict(self) -> dict:
+        return {
+            "text": self.text,
+            "threshold": self.threshold,
+            "n_flagged": len(self.flagged),
+            "entities": [{**asdict(e), "risk": round(e.risk, 4)} for e in self.entities],
+        }
+
+    def __str__(self) -> str:
+        return self.highlight()
+
+    def _repr_html_(self) -> str:  # notebooks show the highlighted text and the summary
+        return (
+            f"<div style='white-space:pre-wrap;line-height:1.6'>{self.highlight('html')}"
+            f"</div><pre>{html.escape(self.summary())}</pre>"
+        )
+
+
+@dataclass
+class Event:
+    """One step of stream(): new text, and the entities of sentences that just finished."""
+
+    text: str
+    entities: list[Entity] = field(default_factory=list)
+    done: bool = False
+    checked: int = 0  # the text up to this character has been checked
+    ids: list[int] = field(default_factory=list)  # on the last event: every generated token
+
+
+def highlight(text: str, flagged: list[Entity], style: str = "plain") -> str:
+    spans: list[list] = []
+    for e in sorted(flagged, key=lambda e: e.start):
+        if spans and e.start < spans[-1][1]:
+            spans[-1][1], spans[-1][2] = max(spans[-1][1], e.end), max(spans[-1][2], e.risk)
+        else:
+            spans.append([e.start, e.end, e.risk])
+    esc = html.escape if style == "html" else (lambda s: s)
+    out, pos = [], 0
+    for lo, hi, risk in spans:
+        out.append(esc(text[pos:lo]))
+        inner = esc(text[lo:hi])
+        if style == "html":
+            out.append(
+                f'<mark class="hs-risk" style="background:#ffcdd2;color:#000" '
+                f'title="risk {risk:.0%}">{inner}</mark>'
+            )
+        elif style == "ansi":
+            out.append(f"\033[1;97;41m{inner}\033[0m")
+        else:
+            out.append(f"[[{inner}]]")
+        pos = hi
+    out.append(esc(text[pos:]))
+    return "".join(out)
+
+
+def _default_device() -> str:
+    if torch.cuda.is_available():
+        return "cuda"
+    if torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"
+
+
+def _weights_dir(weights: str | Path) -> Path:
+    if str(weights) in BUNDLED:
+        return Path(str(files("hallustream") / "weights" / str(weights)))
+    return Path(weights)
+
+
+def _check_weights(d: Path, weights: str | Path) -> None:
+    if not (d / "model.json").exists():
+        raise FileNotFoundError(
+            f"{weights!r} is not a bundled probe {list(BUNDLED)} or a folder with model.json"
+        )
+
+
+class Detector:
+    """A model, its probe, and the thresholds for flagging an entity."""
+
+    def __init__(self, model, tokenizer, probe: Probe, threshold: str | float = 0.5,
+                 chat_template_kwargs: dict | None = None, thresholds: dict | None = None,
+                 entity_finder: EntityFinder | None = None):  # fmt: skip
+        self.model, self.tokenizer, self.probe = model, tokenizer, probe
+        self.chat = chat_template_kwargs or {}
+        self.thresholds = thresholds or {}
+        self.threshold = self._threshold(threshold)
+        # GLiNER runs on the model's GPU when there is one (its MPS support is uneven)
+        gliner_device = "cuda" if str(getattr(model, "device", "cpu")).startswith("cuda") else "cpu"
+        self.find_entities = entity_finder or partial(_entities.find, device=gliner_device)
+
+    @classmethod
+    def from_pretrained(cls, weights: str | Path = "qwen3-4b", model=None, tokenizer=None,
+                        device: str | None = None, dtype: torch.dtype = torch.bfloat16,
+                        threshold: str | float = "balanced",
+                        entity_finder: EntityFinder | None = None) -> Detector:  # fmt: skip
+        """Load bundled weights by name ("qwen3-4b") or a weights folder, and the model they
+        were trained for (unless model and tokenizer are given)."""
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+
+        d = _weights_dir(weights)
+        _check_weights(d, weights)
+        meta = json.loads((d / "model.json").read_text())
+        if model is None:
+            model = AutoModelForCausalLM.from_pretrained(meta["hf_id"], dtype=dtype)
+            model.to(device or _default_device())  # pyright: ignore[reportArgumentType]
+            model.eval()
+        tokenizer = tokenizer or AutoTokenizer.from_pretrained(meta["hf_id"])
+        seeds = sorted(d.glob("probe_seed*.safetensors"))
+        if not seeds:
+            raise FileNotFoundError(f"no probe weights (probe_seed*.safetensors) in {d}")
+        probe = Probe(seeds, d / ADAPTER, device=str(model.device))
+        if probe.weight.shape[0] != model.config.hidden_size:
+            raise ValueError(f"probe width {probe.weight.shape[0]} does not match the model")
+        th = d / "thresholds.json"
+        return cls(model, tokenizer, probe, threshold, meta.get("chat_template_kwargs"),
+                   json.loads(th.read_text()) if th.exists() else {}, entity_finder)  # fmt: skip
+
+    def _threshold(self, threshold: str | float | None) -> float:
+        if threshold is None:
+            return self.threshold
+        if isinstance(threshold, str):
+            presets = [k for k, v in self.thresholds.items() if isinstance(v, dict)]
+            if threshold not in presets:
+                hint = f"use one of {presets}" if presets else "no presets loaded; give a number"
+                raise ValueError(f"unknown threshold {threshold!r}: {hint}")
+            return float(self.thresholds[threshold]["threshold"])
+        if not 0 <= float(threshold) <= 1:
+            raise ValueError(f"threshold {threshold} is not a risk between 0 and 1")
+        return float(threshold)
+
+    def _messages(self, prompt: str | None, messages: list[dict] | None) -> list[dict]:
+        if messages is None:
+            if prompt is None:
+                raise ValueError("give a prompt or messages")
+            messages = [{"role": "user", "content": prompt}]
+        return messages
+
+    def _find(self, text: str, sents: list[tuple[int, int]]) -> list[tuple[int, int]]:
+        """Entities to score: none in sentences where the model says it does not know, and
+        never a bare pronoun."""
+        sents = [(lo, hi) for lo, hi in sents if not is_refusal(text[lo:hi])]
+        spans = self.find_entities(text, sents) if sents else []
+        return [(lo, hi) for lo, hi in spans if text[lo:hi].strip().lower() not in PRONOUNS]
+
+    def _entities(self, text: str, fw: Forward, spans, threshold: float) -> list[Entity]:
+        keep, tok_spans = [], []
+        for lo, hi in sorted(set(spans)):
+            if (t := char_to_tokens(fw, lo, hi)) is not None:
+                keep.append((lo, hi))
+                tok_spans.append(t)
+        _, adapted = self.probe.span_scores(fw, tok_spans)
+        risk = [1 / (1 + math.exp(-float(z))) for z in adapted]
+        return [Entity(text[lo:hi], lo, hi, r, r >= threshold)
+                for (lo, hi), r in zip(keep, risk, strict=True)]  # fmt: skip
+
+    def score(self, prompt: str | None = None, text: str = "", *, messages: list[dict] | None
+              = None, spans: list[tuple[int, int]] | None = None,
+              threshold: str | float | None = None) -> Report:  # fmt: skip
+        """Risk of every entity in `text`, an answer the model wrote to `prompt` (or after
+        `messages`). Entities are found automatically unless `spans` are given."""
+        th = self._threshold(threshold)
+        fw = forward(self.model, self.tokenizer, self.probe, self._messages(prompt, messages),
+                     text, **self.chat)  # fmt: skip
+        if spans is None:
+            spans = self._find(text, sentences(text))
+        return Report(text, self._entities(text, fw, spans, th), th)
+
+    def generate(self, prompt: str | None = None, *, messages: list[dict] | None = None,
+                 max_new_tokens: int = 512, threshold: str | float | None = None,
+                 temperature: float = 0.0) -> Report:  # fmt: skip
+        """Write an answer and check it: stream() collected into one Report."""
+        text, found = "", []
+        for ev in self.stream(prompt, messages=messages, max_new_tokens=max_new_tokens,
+                              threshold=threshold, temperature=temperature):  # fmt: skip
+            text += ev.text
+            found += ev.entities
+        return Report(text, sorted(found, key=lambda e: e.start), self._threshold(threshold))
+
+    @torch.no_grad()
+    def stream(self, prompt: str | None = None, *, messages: list[dict] | None = None,
+               max_new_tokens: int = 512, threshold: str | float | None = None,
+               temperature: float = 0.0) -> Iterator[Event]:  # fmt: skip
+        """Generate an answer token by token; when a sentence ends, its entities are scored
+        from the hidden states already computed for generation (no second pass)."""
+        from transformers import DynamicCache
+
+        th = self._threshold(threshold)
+        tok, lm = self.tokenizer, self.model
+        msgs = self._messages(prompt, messages)
+        prefix = tok.apply_chat_template(
+            msgs, add_generation_prompt=True, tokenize=False, **self.chat
+        )
+        ids = tok(prefix, add_special_tokens=False, return_tensors="pt").input_ids.to(lm.device)
+        eos = lm.generation_config.eos_token_id
+        stop = set(eos if isinstance(eos, list) else [eos]) | {tok.eos_token_id}
+        cache = DynamicCache()
+        out = lm(ids, past_key_values=cache, use_cache=True)
+        gen: list[int] = []
+        dec = [0]  # dec[i]: length of the text decoded from the first i tokens
+        hidden: list[np.ndarray] = []
+        probs: list[np.ndarray] = []
+        text, scored, seen = "", set(), set()  # sentences and entity spans already scored
+        for _ in range(max_new_tokens):
+            logits = out.logits[0, -1].float()
+            if temperature > 0:
+                nxt = int(torch.multinomial(torch.softmax(logits / temperature, -1), 1))
+            else:
+                nxt = int(logits.argmax())
+            if nxt in stop:
+                break
+            gen.append(nxt)
+            out = lm(torch.tensor([[nxt]], device=lm.device), past_key_values=cache,
+                     use_cache=True, output_hidden_states=True)  # fmt: skip
+            h = out.hidden_states[self.probe.layer + 1][0, -1]
+            hidden.append(h.half().cpu().numpy())
+            probs.append(self.probe.token_probs(h[None])[0])
+            ended = text.rstrip()[-1:] in SENTENCE_END or "\n" in text[-1:]
+            new = tok.decode(gen, skip_special_tokens=True)
+            dec.append(len(new))
+            delta = new[len(text) :] if new.startswith(text) else new
+            text = new
+            found = []
+            if ended:
+                found = self._finished(text, gen, dec, hidden, probs, scored, seen, th, False)
+            yield Event(delta, found, checked=max((hi for _, hi in scored), default=0))
+        last = self._finished(text, gen, dec, hidden, probs, scored, seen, th, True)
+        yield Event("", last, True, len(text), gen)
+
+    def _finished(self, text, gen, dec, hidden, probs, scored, seen, th, final) -> list[Entity]:
+        """Entities of sentences that are complete and not scored yet (each span once, even if
+        sentence boundaries shift as more text arrives)."""
+        sents = sentences(text)
+        todo = [s for s in (sents if final else sents[:-1]) if s not in scored]
+        if not todo or not gen:
+            return []
+        scored.update(todo)
+        n = len(gen)
+        fw = Forward(gen, list(zip(dec[:-1], dec[1:], strict=True)), 0, np.stack(probs),
+                     np.stack(hidden), np.full(n, np.nan), np.full(n, np.nan))  # fmt: skip
+        spans = [s for s in self._find(text, todo) if s not in seen]
+        seen.update(spans)
+        return self._entities(text, fw, spans, th)

hallustream/entities.py

@@ -0,0 +1,75 @@
+"""Find the entities to score: names, places, dates, numbers, titles. Uses GLiNER, chosen on dev
+biographies for finding the most false entities; install with
+`pip install "hallustream[entities]"`, or pass your own spans."""
+
+from __future__ import annotations
+
+from functools import lru_cache
+
+# Entity types the judge labels (EVAL_PROTOCOL.md): people, organizations, places, dates,
+# quantities, titles of works, awards; plus the descriptive roles it often includes.
+GLINER_LABELS = [
+    "person", "organization", "location", "date", "number", "nationality",
+    "occupation", "title of work", "award", "event",
+]  # fmt: skip
+GLINER_MODEL = "urchade/gliner_medium-v2.1"  # Apache-2.0
+# Generic lowercase words ("region", "book", "team") are skipped: the fact checker almost never
+# labels them, and on the library-check texts they drew 16-19% of all flags while dropping them
+# kept 96-99% of the false entities found (research/entity_rule.py). Jobs and roles ("actress")
+# stay: in biographies they are claims, and the probe ranks them well (0.92 AUROC).
+ROLES = {"occupation"}
+TITLES = {"mr", "mrs", "ms", "miss", "dr", "prof", "sir", "dame", "st"}  # alone, not a detail
+# Words that are capitalized only because they start a sentence ("The death of ...")
+STARTERS = {
+    "the", "a", "an", "his", "her", "their", "its", "this", "that", "these", "those",
+    "in", "on", "at", "after", "before", "during", "by", "for", "with", "from", "as",
+}  # fmt: skip
+
+
+def available() -> bool:
+    try:
+        import gliner  # noqa: F401  # pyright: ignore[reportMissingImports]
+    except ImportError:
+        return False
+    return True
+
+
+def is_specific(text: str, label: str, sentence_start: bool = False) -> bool:
+    """A name, place, title, date or number (it has a capital letter or a digit), or a role.
+
+    At the start of a sentence, a capitalized function word ("The death of ...") doesn't
+    count, and a lone title ("Mr.") is skipped (EVAL_PROTOCOL.md, follow-ups C)."""
+    words = text.split()
+    if not words or text.strip().rstrip(".").lower() in TITLES:
+        return False
+    if label in ROLES:
+        return True
+    if sentence_start and words[0].lower() in STARTERS:
+        text = " ".join(words[1:])
+    return any(c.isupper() or c.isdigit() for c in text)
+
+
+@lru_cache(maxsize=1)
+def _gliner(device: str = "cpu"):
+    from gliner import GLiNER  # pyright: ignore[reportMissingImports]
+
+    return GLiNER.from_pretrained(GLINER_MODEL).to(device)
+
+
+def find(text: str, sents: list[tuple[int, int]], threshold: float = 0.4,
+         device: str = "cpu") -> list[tuple]:  # fmt: skip
+    """(start, end) character spans of specific entities, found sentence by sentence."""
+    if not available():
+        raise ImportError(
+            "finding entities needs GLiNER: pip install 'hallustream[entities]',"
+            " or pass spans=[(start, end), ...] yourself"
+        )
+    model = _gliner(device)
+    out = []
+    for lo, hi in sents:
+        for e in model.predict_entities(text[lo:hi], GLINER_LABELS, threshold=threshold):
+            a, b = lo + e["start"], lo + e["end"]
+            start = not text[lo:a].strip()  # nothing but spaces before it in the sentence
+            if is_specific(text[a:b], e["label"], start):
+                out.append((a, b))
+    return out