mirror of
https://github.com/denizsafak/abogen.git
synced 2026-07-18 21:50:28 +02:00
445 lines
14 KiB
Python
445 lines
14 KiB
Python
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import os
|
||
import re
|
||
import threading
|
||
import time
|
||
from collections import Counter
|
||
from dataclasses import dataclass, field
|
||
from typing import Any, Dict, Iterable, List, Mapping, Optional, Tuple
|
||
|
||
try: # pragma: no cover - fallback when spaCy not available during tests
|
||
import spacy # type: ignore[import-not-found]
|
||
except Exception: # pragma: no cover - spaCy optional during runtime bootstrap
|
||
spacy = None
|
||
|
||
_Language = Any # type: ignore[misc,assignment]
|
||
Doc = Any # type: ignore[misc,assignment]
|
||
Span = Any # type: ignore[misc,assignment]
|
||
|
||
|
||
_TITLE_PREFIXES = (
|
||
"mr",
|
||
"mrs",
|
||
"ms",
|
||
"miss",
|
||
"dr",
|
||
"prof",
|
||
"sir",
|
||
"madam",
|
||
"lady",
|
||
"lord",
|
||
"capt",
|
||
"captain",
|
||
"col",
|
||
"colonel",
|
||
"maj",
|
||
"major",
|
||
"sgt",
|
||
"sergeant",
|
||
"rev",
|
||
"father",
|
||
"mother",
|
||
"brother",
|
||
"sister",
|
||
)
|
||
|
||
_STOP_LABELS = {
|
||
"the",
|
||
"that",
|
||
"this",
|
||
"those",
|
||
"these",
|
||
"there",
|
||
"here",
|
||
"then",
|
||
"and",
|
||
"but",
|
||
"or",
|
||
"nor",
|
||
"so",
|
||
"yet",
|
||
"dr",
|
||
"mr",
|
||
"mrs",
|
||
"ms",
|
||
"miss",
|
||
"sir",
|
||
"madam",
|
||
"lady",
|
||
"lord",
|
||
}
|
||
|
||
_EXCLUDED_NER_LABELS = {
|
||
"CARDINAL",
|
||
"DATE",
|
||
"ORDINAL",
|
||
"PERCENT",
|
||
"TIME",
|
||
"LAW",
|
||
"MONEY",
|
||
"QUANTITY",
|
||
}
|
||
|
||
_TITLE_PATTERN = re.compile(r"^(?:" + "|".join(re.escape(prefix) for prefix in _TITLE_PREFIXES) + r")\.?\s+", re.IGNORECASE)
|
||
_POSSESSIVE_PATTERN = re.compile(r"(?:'s|’s|\u2019s)$", re.IGNORECASE)
|
||
_NON_WORD_PATTERN = re.compile(r"[^\w\s'-]+")
|
||
_MULTI_SPACE_PATTERN = re.compile(r"\s+")
|
||
_SUFFIX_PATTERN = re.compile(
|
||
r",?\s+(?:jr|sr|ii|iii|iv|v|vi|md|phd|esq|esquire|dds|dvm)\.?$",
|
||
re.IGNORECASE,
|
||
)
|
||
|
||
|
||
@dataclass(slots=True)
|
||
class EntityRecord:
|
||
key: Tuple[str, str]
|
||
label: str
|
||
kind: str
|
||
category: str
|
||
count: int = 0
|
||
samples: List[Dict[str, Any]] = field(default_factory=list)
|
||
chapter_indices: set[int] = field(default_factory=set)
|
||
forms: Counter = field(default_factory=Counter)
|
||
first_position: Optional[Tuple[int, int]] = None
|
||
|
||
def register(self, *, chapter_index: int, position: int, text: str, sentence: Optional[str]) -> None:
|
||
self.count += 1
|
||
self.chapter_indices.add(chapter_index)
|
||
self.forms[text] += 1
|
||
if self.first_position is None:
|
||
self.first_position = (chapter_index, position)
|
||
if sentence and len(self.samples) < 5:
|
||
payload = {
|
||
"excerpt": sentence.strip(),
|
||
"chapter_index": chapter_index,
|
||
}
|
||
if payload not in self.samples:
|
||
self.samples.append(payload)
|
||
|
||
def as_dict(self, ordinal: int) -> Dict[str, Any]:
|
||
chapter_indices = sorted(self.chapter_indices)
|
||
first_chapter = chapter_indices[0] if chapter_indices else None
|
||
return {
|
||
"id": f"{self.category}_{ordinal}",
|
||
"label": self.label,
|
||
"normalized": self.key[1],
|
||
"category": self.category,
|
||
"kind": self.kind,
|
||
"count": self.count,
|
||
"samples": list(self.samples),
|
||
"chapter_indices": chapter_indices,
|
||
"first_chapter": first_chapter,
|
||
"forms": self.forms.most_common(6),
|
||
}
|
||
|
||
|
||
@dataclass(slots=True)
|
||
class EntityExtractionResult:
|
||
summary: Dict[str, Any]
|
||
cache_key: str
|
||
elapsed: float
|
||
errors: List[str]
|
||
|
||
|
||
class EntityModelError(RuntimeError):
|
||
pass
|
||
|
||
|
||
_MODEL_CACHE: Dict[str, Any] = {}
|
||
_MODEL_LOCK = threading.RLock()
|
||
|
||
|
||
def _resolve_model_name(language: str) -> str:
|
||
override = os.environ.get("ABOGEN_SPACY_MODEL")
|
||
if override:
|
||
return override.strip()
|
||
lowered = language.strip().lower()
|
||
if lowered.startswith("en"):
|
||
return "en_core_web_sm"
|
||
return "en_core_web_sm"
|
||
|
||
|
||
def _load_model(language: str) -> Any:
|
||
if spacy is None:
|
||
raise EntityModelError("spaCy is not available. Install spaCy to enable entity extraction.")
|
||
|
||
model_name = _resolve_model_name(language)
|
||
cache_key = model_name.lower()
|
||
with _MODEL_LOCK:
|
||
if cache_key in _MODEL_CACHE:
|
||
return _MODEL_CACHE[cache_key]
|
||
try:
|
||
nlp = spacy.load(model_name) # type: ignore[arg-type]
|
||
except OSError as exc: # pragma: no cover - external dependency failure
|
||
raise EntityModelError(
|
||
f"spaCy model '{model_name}' is not installed. Download it with "
|
||
"`python -m spacy download en_core_web_sm`."
|
||
) from exc
|
||
nlp.max_length = max(nlp.max_length, 2_000_000)
|
||
_MODEL_CACHE[cache_key] = nlp
|
||
return nlp
|
||
|
||
|
||
def _normalize_label(text: str) -> str:
|
||
if not text:
|
||
return ""
|
||
stripped = text.strip().strip("\"'`“”’")
|
||
if not stripped:
|
||
return ""
|
||
stripped = _TITLE_PATTERN.sub("", stripped)
|
||
stripped = _SUFFIX_PATTERN.sub("", stripped)
|
||
stripped = _POSSESSIVE_PATTERN.sub("", stripped)
|
||
stripped = _NON_WORD_PATTERN.sub(" ", stripped)
|
||
stripped = _MULTI_SPACE_PATTERN.sub(" ", stripped)
|
||
stripped = stripped.strip()
|
||
if not stripped or stripped.lower() in _STOP_LABELS:
|
||
return ""
|
||
parts = stripped.split()
|
||
if not parts:
|
||
return ""
|
||
if len(parts) == 1 and len(parts[0]) <= 1:
|
||
return ""
|
||
# Normalise casing: preserve uppercase abbreviations, otherwise title case.
|
||
normalized_parts = []
|
||
for index, part in enumerate(parts):
|
||
if part.isupper():
|
||
normalized_parts.append(part)
|
||
elif part[:1].isupper():
|
||
normalized_parts.append(part[:1].upper() + part[1:])
|
||
elif index == 0:
|
||
normalized_parts.append(part[:1].upper() + part[1:])
|
||
else:
|
||
normalized_parts.append(part)
|
||
normalized = " ".join(normalized_parts).strip()
|
||
if normalized.lower() in _STOP_LABELS:
|
||
return ""
|
||
return normalized
|
||
|
||
|
||
def _token_key(value: str) -> str:
|
||
return _MULTI_SPACE_PATTERN.sub(" ", value.lower().strip()).strip()
|
||
|
||
|
||
def _iter_named_entities(doc: Any) -> Iterable[Any]: # type: ignore[override]
|
||
for ent in getattr(doc, "ents", ()):
|
||
if ent.label_ == "":
|
||
continue
|
||
yield ent
|
||
|
||
|
||
def _extract_propn_tokens(doc: Any) -> Iterable[Any]: # type: ignore[override]
|
||
seen: set[Tuple[int, int]] = set()
|
||
for ent in getattr(doc, "ents", ()): # guard multi-token spans
|
||
seen.add((ent.start, ent.end))
|
||
for token in doc:
|
||
if token.pos_ != "PROPN":
|
||
continue
|
||
span_key = (token.i, token.i + 1)
|
||
if span_key in seen:
|
||
continue
|
||
if token.is_stop:
|
||
continue
|
||
text = token.text.strip()
|
||
if not text:
|
||
continue
|
||
if token.ent_type_:
|
||
continue
|
||
yield doc[token.i : token.i + 1]
|
||
|
||
|
||
def _empty_result(cache_key: str, error: Optional[str] = None) -> EntityExtractionResult:
|
||
payload = {
|
||
"people": [],
|
||
"entities": [],
|
||
"index": {"tokens": []},
|
||
"stats": {
|
||
"tokens": 0,
|
||
"chapters": 0,
|
||
"processed": False,
|
||
},
|
||
"model": None,
|
||
}
|
||
errors = [error] if error else []
|
||
return EntityExtractionResult(summary=payload, cache_key=cache_key, elapsed=0.0, errors=errors)
|
||
|
||
|
||
def extract_entities(
|
||
chapters: Iterable[Mapping[str, Any]],
|
||
*,
|
||
language: str = "en",
|
||
) -> EntityExtractionResult:
|
||
start = time.perf_counter()
|
||
normalized_language = language or "en"
|
||
combined_hasher = hashlib.sha1()
|
||
chapter_texts: List[Tuple[int, str]] = []
|
||
for idx, chapter in enumerate(chapters):
|
||
text = chapter.get("text") if isinstance(chapter, Mapping) else None
|
||
text_value = str(text or "")
|
||
original_index = idx
|
||
if isinstance(chapter, Mapping):
|
||
try:
|
||
original_index = int(chapter.get("index", idx))
|
||
except (TypeError, ValueError):
|
||
original_index = idx
|
||
chapter_texts.append((original_index, text_value))
|
||
if text_value:
|
||
combined_hasher.update(text_value.encode("utf-8", "ignore"))
|
||
combined_hasher.update(str(original_index).encode("utf-8", "ignore"))
|
||
cache_key = combined_hasher.hexdigest()
|
||
|
||
if not chapter_texts:
|
||
return _empty_result(cache_key)
|
||
|
||
try:
|
||
nlp = _load_model(normalized_language)
|
||
except EntityModelError as exc:
|
||
return _empty_result(cache_key, str(exc))
|
||
|
||
records: Dict[Tuple[str, str], EntityRecord] = {}
|
||
tokens_for_index: Dict[str, Dict[str, Any]] = {}
|
||
processed_tokens = 0
|
||
|
||
for chapter_index, text in chapter_texts:
|
||
trimmed = text.strip()
|
||
if not trimmed:
|
||
continue
|
||
if len(trimmed) + 1024 > nlp.max_length:
|
||
nlp.max_length = len(trimmed) + 1024
|
||
doc = nlp(trimmed)
|
||
|
||
def _register_span(span: Any, category_hint: Optional[str] = None) -> None:
|
||
nonlocal processed_tokens
|
||
if category_hint is None and span.label_ in _EXCLUDED_NER_LABELS:
|
||
return
|
||
cleaned = _normalize_label(span.text)
|
||
if not cleaned:
|
||
return
|
||
key = _token_key(cleaned)
|
||
if not key:
|
||
return
|
||
category = category_hint or ("people" if span.label_ == "PERSON" else "entities")
|
||
record_key = (category, key)
|
||
record = records.get(record_key)
|
||
if record is None:
|
||
record = EntityRecord(
|
||
key=record_key,
|
||
label=cleaned,
|
||
kind=span.label_ or ("PROPN" if category == "entities" else "PERSON"),
|
||
category=category,
|
||
)
|
||
records[record_key] = record
|
||
sentence = span.sent.text if hasattr(span, "sent") and span.sent is not None else None
|
||
record.register(
|
||
chapter_index=chapter_index,
|
||
position=span.start,
|
||
text=span.text,
|
||
sentence=sentence,
|
||
)
|
||
processed_tokens += 1
|
||
index_entry = tokens_for_index.get(key)
|
||
if index_entry is None:
|
||
index_entry = {
|
||
"token": record.label,
|
||
"normalized": key,
|
||
"category": category,
|
||
"count": 0,
|
||
"samples": [],
|
||
}
|
||
tokens_for_index[key] = index_entry
|
||
index_entry["count"] += 1
|
||
if sentence and len(index_entry["samples"]) < 3:
|
||
if sentence not in index_entry["samples"]:
|
||
index_entry["samples"].append(sentence)
|
||
|
||
for ent in _iter_named_entities(doc):
|
||
_register_span(ent)
|
||
|
||
for span in _extract_propn_tokens(doc):
|
||
_register_span(span, category_hint="entities")
|
||
|
||
elapsed = time.perf_counter() - start
|
||
|
||
people_records = [record for record in records.values() if record.category == "people"]
|
||
people_keys = {record.key[1] for record in people_records}
|
||
entity_records = [
|
||
record
|
||
for record in records.values()
|
||
if record.category == "entities"
|
||
and record.key[1] not in people_keys
|
||
and record.kind != "PERSON"
|
||
]
|
||
|
||
people_records.sort(key=lambda rec: (-rec.count, rec.label))
|
||
entity_records.sort(key=lambda rec: (-rec.count, rec.label))
|
||
|
||
people_payload = [record.as_dict(index + 1) for index, record in enumerate(people_records)]
|
||
entity_payload = [record.as_dict(index + 1) for index, record in enumerate(entity_records)]
|
||
|
||
index_payload = sorted(tokens_for_index.values(), key=lambda item: (-item["count"], item["token"]))
|
||
|
||
summary = {
|
||
"people": people_payload,
|
||
"entities": entity_payload,
|
||
"index": {"tokens": index_payload},
|
||
"stats": {
|
||
"tokens": processed_tokens,
|
||
"chapters": len(chapter_texts),
|
||
"processed": True,
|
||
"people": len(people_payload),
|
||
"entities": len(entity_payload),
|
||
},
|
||
"model": {
|
||
"name": getattr(nlp, "meta", {}).get("name", "unknown"),
|
||
"version": getattr(nlp, "meta", {}).get("version", "unknown"),
|
||
"lang": getattr(nlp, "meta", {}).get("lang", normalized_language),
|
||
},
|
||
}
|
||
|
||
return EntityExtractionResult(summary=summary, cache_key=cache_key, elapsed=elapsed, errors=[])
|
||
|
||
|
||
def search_tokens(index: Mapping[str, Any], query: str, *, limit: int = 15) -> List[Dict[str, Any]]:
|
||
tokens = index.get("tokens") if isinstance(index, Mapping) else None
|
||
if not isinstance(tokens, list) or not query:
|
||
return []
|
||
normalized = query.strip().lower()
|
||
if not normalized:
|
||
return tokens[:limit]
|
||
results: List[Dict[str, Any]] = []
|
||
for entry in tokens:
|
||
token_label = str(entry.get("token", ""))
|
||
normalized_label = token_label.lower()
|
||
if normalized in normalized_label or normalized in str(entry.get("normalized", "")):
|
||
results.append(entry)
|
||
if len(results) >= limit:
|
||
break
|
||
return results
|
||
|
||
|
||
def merge_override(summary: Mapping[str, Any], overrides: Mapping[str, Mapping[str, Any]]) -> Dict[str, Any]:
|
||
if not isinstance(summary, Mapping):
|
||
return {"people": [], "entities": []}
|
||
merged_summary: Dict[str, Any] = dict(summary)
|
||
for key in ("people", "entities"):
|
||
items = summary.get(key)
|
||
if not isinstance(items, list):
|
||
continue
|
||
merged_items: List[Dict[str, Any]] = []
|
||
for entry in items:
|
||
if not isinstance(entry, Mapping):
|
||
continue
|
||
normalized = _token_key(str(entry.get("normalized") or entry.get("label") or ""))
|
||
merged = dict(entry)
|
||
if normalized and normalized in overrides:
|
||
merged_override = dict(overrides[normalized])
|
||
merged["override"] = merged_override
|
||
merged_items.append(merged)
|
||
merged_summary[key] = merged_items
|
||
return merged_summary
|
||
|
||
|
||
def normalize_token(token: str) -> str:
|
||
return _token_key(_normalize_label(token))
|