Coverage for src/lilbee/wiki/entity_extractor/ner_concepts.py: 100%
113 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""spaCy NER entity extractor (default strategy).
3Produces typed NER entities only. LLM-curated concept pages are
4proposed downstream by the per-source batched call in
5:mod:`lilbee.wiki.generation`.
6"""
8from __future__ import annotations
10import functools
11import logging
12import re
13import threading
14from typing import TYPE_CHECKING, Any, NamedTuple
16from lilbee.core.text import collapse_whitespace, is_valid_label, make_slug, strip_possessive
17from lilbee.wiki.entity_extractor.base import (
18 ChunkRef,
19 EntityKind,
20 ExtractedEntity,
21)
23if TYPE_CHECKING:
24 from lilbee.core.config import Config
25 from lilbee.data.store import SearchChunk
26 from lilbee.providers.base import LLMProvider
28log = logging.getLogger(__name__)
31# Pre-spaCy markdown-noise strippers. Compiled once at module scope so
32# the extractor's hot path does not recompile them per chunk. Match on
33# line boundaries via re.MULTILINE; each sub() empties the matched
34# line so downstream line-joins collapse the hole to a single newline.
35_TABLE_ROW_RE = re.compile(r"^\|.*\|\s*$", re.MULTILINE)
36_PAGE_NUMBER_RE = re.compile(r"^\s*\d{1,4}\s*$", re.MULTILINE)
37_NAV_CHROME_RE = re.compile(
38 r"^\s*(?:Home|Menu|Navigation|Edit this page|Jump to navigation|Jump to search)\s*$",
39 re.MULTILINE,
40)
43def pre_clean_for_ner(text: str) -> str:
44 """Strip markdown-structural noise before handing text to spaCy.
46 Removes whole-line markdown-table rows (``| Designer | Irv ... |``),
47 standalone page-number lines from PDF extraction (``42``), and
48 Wikipedia / CMS navigation chrome (``Edit this page``). Leaves
49 prose untouched: every regex anchors to a full line and emits an
50 empty line in place of the match, which spaCy treats as a sentence
51 break.
53 Only targets the noise patterns actually observed in the bb-8b7s
54 QA corpus. Fuller markdown parsing is deferred; a regex pre-clean
55 is sufficient for the current signal-to-noise ratio.
56 """
57 text = _TABLE_ROW_RE.sub("", text)
58 text = _PAGE_NUMBER_RE.sub("", text)
59 return _NAV_CHROME_RE.sub("", text)
62class NerConceptsExtractor:
63 """Emit typed NER entities (``EntityKind.ENTITY`` only).
65 LLM-curated concept pages are produced downstream by the per-source
66 batched call in :mod:`lilbee.wiki.generation`.
67 """
69 def __init__(self, provider: LLMProvider, config: Config) -> None:
70 self._provider = provider
71 self._config = config
73 def available(self, *, quiet: bool = False) -> bool:
74 """True when the spaCy pipeline is loadable. Cached, so this is cheap."""
75 return (_shared_spacy().pipeline if quiet else _load_spacy()) is not None
77 def extract(self, chunks: list[SearchChunk]) -> list[ExtractedEntity]:
78 if not chunks:
79 return []
80 nlp = _load_spacy()
81 if nlp is None:
82 return []
84 entity_records: dict[str, _Aggregate] = {}
85 allowed_ent_types = self._config.concept_allowed_ent_types
87 debug_enabled = log.isEnabledFor(logging.DEBUG)
88 funnel: dict[str, int] = {
89 "raw_ents": 0,
90 "type_filter_dropped": 0,
91 "label_sanity_dropped_entities": 0,
92 "kept_entity_surfaces": 0,
93 }
94 cleaned_texts = (pre_clean_for_ner(c.chunk) for c in chunks)
95 # _load_spacy returns a process-global Language shared by every extractor
96 # instance; a spaCy Language is not safe for concurrent processing, so
97 # serialize the whole lazy nlp.pipe iteration behind the module lock.
98 with _NLP_LOCK:
99 for chunk, doc in zip(chunks, nlp.pipe(cleaned_texts), strict=True):
100 ref = ChunkRef(source=chunk.source, chunk_index=chunk.chunk_index)
101 _accumulate_doc_entities(
102 doc, ref, entity_records, allowed_ent_types, funnel, debug_enabled
103 )
105 _merge_possessives(entity_records)
107 if debug_enabled:
108 log.debug(
109 "ner funnel: raw_ents=%(raw_ents)d "
110 "type_filter_dropped=%(type_filter_dropped)d "
111 "label_sanity_dropped_entities=%(label_sanity_dropped_entities)d "
112 "kept_entity_surfaces=%(kept_entity_surfaces)d",
113 funnel,
114 )
116 min_mentions = self._config.wiki_entity_min_mentions
117 results: list[ExtractedEntity] = []
118 for agg in entity_records.values():
119 record = _make_record(agg, EntityKind.ENTITY, min_mentions)
120 if record is not None:
121 results.append(record)
122 results.sort(key=lambda e: (e.kind.value, e.slug))
123 return results
126def _accumulate_doc_entities(
127 doc: Any,
128 ref: ChunkRef,
129 entity_records: dict[str, _Aggregate],
130 allowed_ent_types: set[str] | frozenset[str],
131 funnel: dict[str, int],
132 debug_enabled: bool,
133) -> None:
134 """Fold one spaCy doc's entities into ``entity_records`` (mutated in place)."""
135 for ent in doc.ents:
136 funnel["raw_ents"] += 1
137 if ent.label_ not in allowed_ent_types:
138 funnel["type_filter_dropped"] += 1
139 continue
140 # Collapse first: spaCy spans cross wrapped lines constantly in PDF
141 # text, and the label is interpolated raw into a single-line marker.
142 # Rejecting the wrapped form instead would lose the mention entirely.
143 surface = collapse_whitespace(ent.text)
144 if not is_valid_label(surface):
145 funnel["label_sanity_dropped_entities"] += 1
146 if debug_enabled:
147 log.debug("label-sanity: rejected entity %r", surface)
148 continue
149 # Key by slug, not the normalized surface: the index keys stubs on
150 # make_slug, so two surfaces that slug alike ("Ford Motor Co." and
151 # "Ford Motor Co") are one subject. Keyed by surface they stay two
152 # aggregates that an incremental refresh sums, double-counting mentions
153 # and duplicating refs past the floor.
154 key = make_slug(surface)
155 rec = entity_records.setdefault(key, _Aggregate(label=surface, type_hint=ent.label_))
156 rec.refs.add(ref)
157 funnel["kept_entity_surfaces"] += 1
160def _merge_possessives(entity_records: dict[str, _Aggregate]) -> None:
161 """Fold each possessive surface into its standalone base entity, in place.
163 "Solar System's" and "Solar System" are one subject, but a name whose
164 canonical form carries the clitic (McDonald's) has no standalone base in
165 the corpus, so a possessive with no base aggregate stays as it is.
166 """
167 for key in list(entity_records):
168 agg = entity_records[key]
169 base = strip_possessive(agg.label)
170 if base == agg.label:
171 continue
172 base_agg = entity_records.get(make_slug(base))
173 if base_agg is None or base_agg is agg:
174 continue
175 base_agg.refs |= agg.refs
176 del entity_records[key]
179class _Aggregate:
180 """Mutable accumulator used only while folding per-chunk hits."""
182 __slots__ = ("label", "refs", "type_hint")
184 def __init__(self, label: str, type_hint: str) -> None:
185 self.label = label
186 self.type_hint = type_hint
187 self.refs: set[ChunkRef] = set()
190def _sorted_refs(refs: set[ChunkRef]) -> tuple[ChunkRef, ...]:
191 return tuple(sorted(refs, key=lambda r: (r.source, r.chunk_index)))
194def _make_record(agg: _Aggregate, kind: EntityKind, min_mentions: int) -> ExtractedEntity | None:
195 """Turn an aggregate into an ``ExtractedEntity`` or drop it.
197 Filters records below the mention threshold and records whose label
198 slug-cleans to an empty string (e.g. labels of only punctuation);
199 without the empty-slug guard those would try to write files named
200 just ``.md`` on disk.
201 """
202 if len(agg.refs) < min_mentions:
203 return None
204 slug = make_slug(agg.label)
205 if not slug:
206 return None
207 return ExtractedEntity(
208 slug=slug,
209 kind=kind,
210 label=agg.label,
211 type_hint=agg.type_hint,
212 chunk_refs=_sorted_refs(agg.refs),
213 )
216_NLP_LOCK = threading.Lock()
217"""Serializes use of the shared (``@functools.cache``d) spaCy Language.
219A spaCy Language is not safe for concurrent processing; every extract() call
220runs ``nlp.pipe`` on the same cached instance, so the daemon serializes them."""
223class _Spacy(NamedTuple):
224 """The shared spaCy pipeline, or the name of what its load could not import."""
226 pipeline: Any | None
227 missing: str = ""
230@functools.cache
231def _shared_spacy() -> _Spacy:
232 """Load the shared spaCy pipeline once for the process. Logs nothing."""
233 try:
234 from lilbee.retrieval.concepts import load_spacy_pipeline
235 except ImportError:
236 return _Spacy(None, "lilbee.concepts")
237 try:
238 return _Spacy(load_spacy_pipeline())
239 except ImportError:
240 return _Spacy(None, "spaCy model")
243@functools.cache
244def _warn_disabled(missing: str) -> None:
245 """Warn that extraction is disabled for want of *missing*, once for each process.
247 Without the cache every chunk-extract call repeats the warning: 1 000
248 identical lines on a corpus of 1 000 chunks.
249 """
250 log.warning("Entity extraction disabled: %s unavailable", missing)
253def _load_spacy() -> Any | None:
254 """The shared spaCy pipeline, or None with the warning when it is unavailable."""
255 loaded = _shared_spacy()
256 if loaded.pipeline is None:
257 _warn_disabled(loaded.missing)
258 return loaded.pipeline