Richie
6c8a4bfea7
Refactor protected phrase handling from a single module into a python/ebook_search/protected_phrases package covering extraction, storage, and runtime matching. Phrase filtering is now data-driven via bundled TOML files: ignored_phrases, bad_starts, bad_ends, and most_common_words. Add phrase-tuning settings to EbookSearchConfig so candidate generation, scoring, LLM judging, and matching are configurable rather than hardcoded: token bounds, entity token limit, raw n-gram min count, frequency and chapter-spread score thresholds, candidate/LLM/target caps, confidence threshold, nesting defaults, and the phrase hit boost.
2011 lines
71 KiB
Python
2011 lines
71 KiB
Python
"""Protected phrase extraction, storage, and runtime matching."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib
|
|
import json
|
|
import logging
|
|
import re
|
|
import tomllib
|
|
from collections import defaultdict
|
|
from dataclasses import dataclass
|
|
from functools import cache
|
|
from pathlib import Path
|
|
from time import perf_counter
|
|
from typing import TYPE_CHECKING, Protocol
|
|
|
|
from sqlalchemy import and_, delete, func, or_, select
|
|
|
|
from python.ebook_search.llm_interface import request_chat_completion
|
|
from python.orm.richie import (
|
|
EbookCandidatePhrase,
|
|
EbookChunk,
|
|
EbookChunkPhraseMention,
|
|
EbookPhraseAlias,
|
|
EbookProtectedPhrase,
|
|
EbookSource,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Iterable, Iterator, Mapping, Sequence
|
|
from types import ModuleType
|
|
|
|
from sqlalchemy.orm import Session
|
|
|
|
from python.ebook_search.config import EbookSearchConfig
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
PHRASE_CONFIG_DIR = Path(__file__).resolve().parent
|
|
BAD_START_SCORE_PENALTY = 3.0
|
|
BAD_END_SCORE_PENALTY = 3.0
|
|
SOURCE_FIELDS = (
|
|
"source_raw_ngram",
|
|
"source_yake",
|
|
"source_spacy_ner",
|
|
"source_spacy_noun_chunk",
|
|
"source_capitalized",
|
|
"source_metadata",
|
|
)
|
|
JSON_OBJECT_RE = re.compile(r"\{.*\}", re.DOTALL)
|
|
CAPITALIZED_PHRASE_RE = re.compile(r"\b(?:[A-Z][a-zA-Z']+)(?:\s+(?:of|the|and|in|on|for|[A-Z][a-zA-Z']+)){0,6}")
|
|
|
|
|
|
class SpacySpan(Protocol):
|
|
"""Small protocol for the spaCy span attributes used by this module."""
|
|
|
|
text: str
|
|
|
|
|
|
class SpacyEntity(SpacySpan, Protocol):
|
|
"""Small protocol for the spaCy entity attributes used by this module."""
|
|
|
|
label_: str
|
|
|
|
|
|
class SpacyDoc(Protocol):
|
|
"""Small protocol for the spaCy doc attributes used by this module."""
|
|
|
|
ents: Iterable[SpacyEntity]
|
|
noun_chunks: Iterable[SpacySpan]
|
|
|
|
|
|
class SpacyLanguage(Protocol):
|
|
"""Small protocol for a callable spaCy language pipeline."""
|
|
|
|
def __call__(self, text: str) -> SpacyDoc:
|
|
"""Parse text into a spaCy-like doc."""
|
|
|
|
|
|
class YakeExtractor(Protocol):
|
|
"""Small protocol for the YAKE extractor used by this module."""
|
|
|
|
def extract_keywords(self, text: str) -> Iterable[tuple[str, float]]:
|
|
"""Return YAKE keyword tuples."""
|
|
|
|
|
|
class YakeExtractorFactory(Protocol):
|
|
"""Callable constructor protocol for YAKE keyword extractors."""
|
|
|
|
def __call__(self, *, lan: str, n: int, dedupLim: float, top: int) -> YakeExtractor: # noqa: N803
|
|
"""Create a YAKE keyword extractor."""
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class NormalizedToken:
|
|
"""A normalized token plus its source character span."""
|
|
|
|
text: str
|
|
start_char: int
|
|
end_char: int
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class PhraseCandidate:
|
|
"""A phrase candidate with merged extraction-source metadata."""
|
|
|
|
phrase_text: str
|
|
phrase_norm: str
|
|
token_count: int
|
|
source_raw_ngram: bool = False
|
|
source_yake: bool = False
|
|
source_spacy_ner: bool = False
|
|
source_spacy_noun_chunk: bool = False
|
|
source_capitalized: bool = False
|
|
source_metadata: bool = False
|
|
spacy_label: str | None = None
|
|
raw_count: int = 0
|
|
chapter_count: int = 0
|
|
yake_score: float | None = None
|
|
candidate_score: float = 0.0
|
|
sample_contexts: tuple[str, ...] = ()
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class LLMJudgment:
|
|
"""A structured phrase judgment returned by the LLM judge."""
|
|
|
|
keep: bool
|
|
canonical: str | None
|
|
category: str | None
|
|
aliases: tuple[str, ...]
|
|
confidence: float
|
|
importance: float = 0.5
|
|
allow_nested: bool = False
|
|
suppress_children: bool = True
|
|
reason: str | None = None
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class PhraseLookup:
|
|
"""In-memory lookup maps used for constant-time phrase-window checks."""
|
|
|
|
norm_to_phrase_ids: Mapping[str, tuple[int, ...]]
|
|
alias_to_phrase_ids: Mapping[str, tuple[int, ...]]
|
|
min_tokens: int
|
|
max_tokens: int
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class PhraseMatch:
|
|
"""An unhydrated query or chunk phrase match."""
|
|
|
|
phrase_norm: str
|
|
start_token: int
|
|
end_token: int
|
|
token_count: int
|
|
phrase_id: int | None = None
|
|
start_char: int | None = None
|
|
end_char: int | None = None
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class HydratedPhraseMatch:
|
|
"""A phrase match with protected-phrase metadata attached."""
|
|
|
|
phrase_id: int
|
|
matched_norm: str
|
|
phrase_text: str
|
|
phrase_norm: str
|
|
canonical_id: str
|
|
phrase_type: str | None
|
|
token_count: int
|
|
confidence: float
|
|
importance: float
|
|
allow_nested: bool
|
|
suppress_children: bool
|
|
start_token: int
|
|
end_token: int
|
|
start_char: int | None = None
|
|
end_char: int | None = None
|
|
book_id: int | None = None
|
|
series_id: int | None = None
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class ProtectedPhraseBackfillResult:
|
|
"""Summary of protected phrase generation for existing books."""
|
|
|
|
books_seen: int
|
|
books_built: int
|
|
protected_phrases: int
|
|
phrase_mentions: int
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class PhraseCandidateGenerationResult:
|
|
"""Summary of candidate phrase extraction for indexed books."""
|
|
|
|
books_seen: int
|
|
books_built: int
|
|
candidate_phrases: int
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class PhraseJudgmentBackfillResult:
|
|
"""Summary of LLM judging for stored candidate phrases."""
|
|
|
|
books_seen: int
|
|
books_judged: int
|
|
candidates_judged: int
|
|
protected_phrases: int
|
|
phrase_mentions: int
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class PhraseRecalculationResult:
|
|
"""Summary of phrase cleanup and candidate regeneration for one book."""
|
|
|
|
book_id: int
|
|
deleted_candidates: int
|
|
deleted_protected_phrases: int
|
|
deleted_aliases: int
|
|
deleted_mentions: int
|
|
candidate_phrases: int
|
|
|
|
|
|
@cache
|
|
def get_ignored_phrases() -> frozenset[str]:
|
|
"""Return ignored phrase strings loaded from TOML."""
|
|
return load_toml_string_set(PHRASE_CONFIG_DIR / "ignored_phrases.toml", "phrases")
|
|
|
|
|
|
@cache
|
|
def get_bad_ends() -> frozenset[str]:
|
|
"""Return bad phrase-ending tokens loaded from TOML."""
|
|
return load_toml_string_set(PHRASE_CONFIG_DIR / "bad_ends.toml", "tokens")
|
|
|
|
|
|
@cache
|
|
def get_bad_starts() -> frozenset[str]:
|
|
"""Return bad phrase-starting tokens loaded from TOML."""
|
|
return load_toml_string_set(PHRASE_CONFIG_DIR / "bad_starts.toml", "tokens")
|
|
|
|
|
|
@cache
|
|
def get_most_common_words() -> frozenset[str]:
|
|
"""Return the most common English words loaded from TOML."""
|
|
return load_toml_string_set(PHRASE_CONFIG_DIR / "most_common_words.toml", "words")
|
|
|
|
|
|
def load_toml_string_set(path: Path, key: str) -> frozenset[str]:
|
|
"""Load and validate a TOML string list as a normalized immutable set."""
|
|
with path.open("rb") as file:
|
|
body = tomllib.load(file)
|
|
|
|
values = body.get(key)
|
|
if not isinstance(values, list) or not all(isinstance(item, str) for item in values):
|
|
msg = f"{path} must contain a {key!r} string list"
|
|
raise ValueError(msg)
|
|
return frozenset(normalize_text(value) for value in values if normalize_text(value))
|
|
|
|
|
|
def normalize_text(text: str) -> str:
|
|
"""Normalize text for phrase storage and lookup."""
|
|
return " ".join(token.text for token in tokenize_with_offsets(text))
|
|
|
|
|
|
def tokenize(text: str) -> list[str]:
|
|
"""Normalize and split text into phrase-detection tokens."""
|
|
return [token.text for token in tokenize_with_offsets(text)]
|
|
|
|
|
|
def tokenize_with_offsets(text: str) -> list[NormalizedToken]:
|
|
"""Normalize text into tokens while preserving original character offsets."""
|
|
tokens: list[NormalizedToken] = []
|
|
current: list[str] = []
|
|
start_char: int | None = None
|
|
|
|
for index, char in enumerate(text):
|
|
normalized = normalize_char(char)
|
|
if normalized == " ":
|
|
if current and start_char is not None:
|
|
tokens.append(NormalizedToken(text="".join(current), start_char=start_char, end_char=index))
|
|
current = []
|
|
start_char = None
|
|
continue
|
|
if start_char is None:
|
|
start_char = index
|
|
current.append(normalized)
|
|
|
|
if current and start_char is not None:
|
|
tokens.append(NormalizedToken(text="".join(current), start_char=start_char, end_char=len(text)))
|
|
return tokens
|
|
|
|
|
|
def normalize_char(char: str) -> str:
|
|
"""Normalize one character into a token character or a separator."""
|
|
if char in {"\u2019", "\u2018"}:
|
|
return "'"
|
|
if char in {"-", "\u2013", "\u2014"}:
|
|
return " "
|
|
|
|
lowered = char.lower()
|
|
if lowered in "abcdefghijklmnopqrstuvwxyz0123456789'":
|
|
return lowered
|
|
return " "
|
|
|
|
|
|
def strip_leading_articles(phrase_norm: str) -> str:
|
|
"""Remove one leading English article from a normalized phrase."""
|
|
tokens_ = phrase_norm.split()
|
|
if tokens_ and tokens_[0] in {"the", "a", "an"}:
|
|
tokens_ = tokens_[1:]
|
|
return " ".join(tokens_)
|
|
|
|
|
|
def normalize_candidate_phrase(
|
|
phrase_text: str,
|
|
config: EbookSearchConfig,
|
|
*,
|
|
min_tokens: int | None = None,
|
|
max_tokens: int | None = None,
|
|
strip_leading_article: bool = False,
|
|
) -> tuple[str, str, int] | None:
|
|
"""Normalize a candidate phrase and validate token bounds."""
|
|
normalized_tokens = tokenize_with_offsets(phrase_text)
|
|
start = 0
|
|
if strip_leading_article and normalized_tokens and normalized_tokens[0].text in {"the", "a", "an"}:
|
|
start = 1
|
|
|
|
selected_tokens = normalized_tokens[start:]
|
|
min_count = config.phrase_min_tokens if min_tokens is None else min_tokens
|
|
max_count = config.phrase_max_tokens if max_tokens is None else max_tokens
|
|
if len(selected_tokens) < min_count or len(selected_tokens) > max_count:
|
|
return None
|
|
|
|
phrase_norm = " ".join(token.text for token in selected_tokens)
|
|
if phrase_norm in get_ignored_phrases():
|
|
return None
|
|
|
|
display_text = phrase_text[selected_tokens[0].start_char : selected_tokens[-1].end_char].strip()
|
|
return display_text or phrase_norm, phrase_norm, len(selected_tokens)
|
|
|
|
|
|
def extract_raw_ngrams(text: str, config: EbookSearchConfig) -> dict[str, PhraseCandidate]:
|
|
"""Extract raw normalized n-grams as high-recall candidates."""
|
|
tokens_ = tokenize(text)
|
|
out: dict[str, PhraseCandidate] = {}
|
|
for ngram_size in range(config.phrase_min_tokens, config.phrase_max_tokens + 1):
|
|
for start in range(len(tokens_) - ngram_size + 1):
|
|
normalized = normalize_candidate_phrase(" ".join(tokens_[start : start + ngram_size]), config)
|
|
if normalized is None:
|
|
continue
|
|
phrase_text, phrase_norm, token_count = normalized
|
|
item = out.setdefault(
|
|
phrase_norm,
|
|
PhraseCandidate(
|
|
phrase_text=phrase_text,
|
|
phrase_norm=phrase_norm,
|
|
token_count=token_count,
|
|
source_raw_ngram=True,
|
|
),
|
|
)
|
|
item.raw_count += 1
|
|
return out
|
|
|
|
|
|
def extract_yake_candidates(
|
|
book_text: str,
|
|
config: EbookSearchConfig,
|
|
top_k: int = 1000,
|
|
) -> dict[str, PhraseCandidate]:
|
|
"""Extract YAKE keyphrases when the optional YAKE package is installed."""
|
|
yake_module = optional_module("yake")
|
|
if yake_module is None:
|
|
logger.warning("ebook_phrase_yake_unavailable_skipping")
|
|
return {}
|
|
|
|
keyword_extractor = yake_module.KeywordExtractor
|
|
extractor = keyword_extractor(lan="en", n=config.phrase_max_tokens, dedupLim=0.85, top=top_k)
|
|
out: dict[str, PhraseCandidate] = {}
|
|
for phrase_text, yake_score in extractor.extract_keywords(book_text):
|
|
normalized = normalize_candidate_phrase(phrase_text, config)
|
|
if normalized is None:
|
|
continue
|
|
display_text, phrase_norm, token_count = normalized
|
|
out[phrase_norm] = PhraseCandidate(
|
|
phrase_text=display_text,
|
|
phrase_norm=phrase_norm,
|
|
token_count=token_count,
|
|
source_yake=True,
|
|
yake_score=float(yake_score),
|
|
)
|
|
return out
|
|
|
|
|
|
def optional_module(module_name: str) -> ModuleType | None:
|
|
"""Import an optional module and return None when it is not installed."""
|
|
try:
|
|
return importlib.import_module(module_name)
|
|
except ImportError:
|
|
return None
|
|
|
|
|
|
def extract_spacy_candidates(
|
|
book_text: str,
|
|
nlp: SpacyLanguage,
|
|
config: EbookSearchConfig,
|
|
) -> dict[str, PhraseCandidate]:
|
|
"""Extract spaCy named entities and noun chunks from one text block."""
|
|
out: dict[str, PhraseCandidate] = {}
|
|
doc = nlp(book_text)
|
|
|
|
for ent in doc.ents:
|
|
normalized = normalize_candidate_phrase(
|
|
ent.text,
|
|
config,
|
|
max_tokens=config.phrase_max_entity_tokens,
|
|
)
|
|
if normalized is None:
|
|
continue
|
|
phrase_text, phrase_norm, token_count = normalized
|
|
out[phrase_norm] = PhraseCandidate(
|
|
phrase_text=phrase_text,
|
|
phrase_norm=phrase_norm,
|
|
token_count=token_count,
|
|
source_spacy_ner=True,
|
|
spacy_label=ent.label_,
|
|
)
|
|
|
|
for chunk in doc.noun_chunks:
|
|
normalized = normalize_candidate_phrase(chunk.text, config, strip_leading_article=True)
|
|
if normalized is None:
|
|
continue
|
|
phrase_text, phrase_norm, token_count = normalized
|
|
out[phrase_norm] = PhraseCandidate(
|
|
phrase_text=phrase_text,
|
|
phrase_norm=phrase_norm,
|
|
token_count=token_count,
|
|
source_spacy_noun_chunk=True,
|
|
)
|
|
return out
|
|
|
|
|
|
def extract_capitalized_phrases(original_text: str, config: EbookSearchConfig) -> dict[str, PhraseCandidate]:
|
|
"""Extract capitalized phrase runs that often carry fictional terms."""
|
|
out: dict[str, PhraseCandidate] = {}
|
|
for match in CAPITALIZED_PHRASE_RE.finditer(original_text):
|
|
phrase_text = match.group(0).strip()
|
|
normalized = normalize_candidate_phrase(
|
|
phrase_text,
|
|
config,
|
|
max_tokens=config.phrase_max_entity_tokens,
|
|
)
|
|
if normalized is None:
|
|
continue
|
|
display_text, phrase_norm, token_count = normalized
|
|
out[phrase_norm] = PhraseCandidate(
|
|
phrase_text=display_text,
|
|
phrase_norm=phrase_norm,
|
|
token_count=token_count,
|
|
source_capitalized=True,
|
|
)
|
|
return out
|
|
|
|
|
|
def extract_metadata_candidates(
|
|
metadata: Mapping[str, object] | None,
|
|
config: EbookSearchConfig,
|
|
) -> dict[str, PhraseCandidate]:
|
|
"""Extract phrases from book metadata values such as title, author, and series."""
|
|
if metadata is None:
|
|
return {}
|
|
|
|
out: dict[str, PhraseCandidate] = {}
|
|
for value in metadata.values():
|
|
if value is None:
|
|
continue
|
|
phrase_text = str(value).strip()
|
|
normalized = normalize_candidate_phrase(
|
|
phrase_text,
|
|
config,
|
|
max_tokens=config.phrase_max_entity_tokens,
|
|
)
|
|
if normalized is None:
|
|
continue
|
|
display_text, phrase_norm, token_count = normalized
|
|
out[phrase_norm] = PhraseCandidate(
|
|
phrase_text=display_text,
|
|
phrase_norm=phrase_norm,
|
|
token_count=token_count,
|
|
source_metadata=True,
|
|
)
|
|
return out
|
|
|
|
|
|
def merge_candidate_sources(*sources: Mapping[str, PhraseCandidate]) -> dict[str, PhraseCandidate]:
|
|
"""Merge candidate dictionaries by normalized phrase."""
|
|
merged: dict[str, PhraseCandidate] = {}
|
|
for source in sources:
|
|
for phrase_norm, item in source.items():
|
|
existing = merged.setdefault(
|
|
phrase_norm,
|
|
PhraseCandidate(
|
|
phrase_text=item.phrase_text,
|
|
phrase_norm=phrase_norm,
|
|
token_count=item.token_count,
|
|
),
|
|
)
|
|
merge_candidate(existing, item)
|
|
return merged
|
|
|
|
|
|
def merge_candidate(existing: PhraseCandidate, item: PhraseCandidate) -> None:
|
|
"""Merge one candidate into an existing candidate object."""
|
|
existing.source_raw_ngram = existing.source_raw_ngram or item.source_raw_ngram
|
|
existing.source_yake = existing.source_yake or item.source_yake
|
|
existing.source_spacy_ner = existing.source_spacy_ner or item.source_spacy_ner
|
|
existing.source_spacy_noun_chunk = existing.source_spacy_noun_chunk or item.source_spacy_noun_chunk
|
|
existing.source_capitalized = existing.source_capitalized or item.source_capitalized
|
|
existing.source_metadata = existing.source_metadata or item.source_metadata
|
|
existing.raw_count += item.raw_count
|
|
if item.yake_score is not None:
|
|
existing.yake_score = item.yake_score
|
|
if item.spacy_label:
|
|
existing.spacy_label = item.spacy_label
|
|
|
|
|
|
def enrich_with_frequency_and_chapter_counts(
|
|
candidates: Mapping[str, PhraseCandidate],
|
|
chapters: Sequence[str],
|
|
) -> dict[str, PhraseCandidate]:
|
|
"""Add raw occurrence and chapter-spread counts to candidates."""
|
|
if not candidates:
|
|
return {}
|
|
|
|
candidate_sets_by_size: dict[int, set[str]] = defaultdict(set)
|
|
for phrase_norm, candidate in candidates.items():
|
|
candidate_sets_by_size[candidate.token_count].add(phrase_norm)
|
|
|
|
total_counts: defaultdict[str, int] = defaultdict(int)
|
|
chapter_counts: defaultdict[str, int] = defaultdict(int)
|
|
for chapter in chapters:
|
|
seen_in_chapter: set[str] = set()
|
|
chapter_tokens = tokenize(chapter)
|
|
for ngram_size, candidate_norms in candidate_sets_by_size.items():
|
|
for start in range(len(chapter_tokens) - ngram_size + 1):
|
|
phrase_norm = " ".join(chapter_tokens[start : start + ngram_size])
|
|
if phrase_norm not in candidate_norms:
|
|
continue
|
|
total_counts[phrase_norm] += 1
|
|
seen_in_chapter.add(phrase_norm)
|
|
for phrase_norm in seen_in_chapter:
|
|
chapter_counts[phrase_norm] += 1
|
|
|
|
enriched = dict(candidates)
|
|
for phrase_norm, candidate in enriched.items():
|
|
candidate.raw_count = max(candidate.raw_count, total_counts[phrase_norm])
|
|
candidate.chapter_count = chapter_counts[phrase_norm]
|
|
return enriched
|
|
|
|
|
|
def filter_storable_candidates(
|
|
candidates: Mapping[str, PhraseCandidate],
|
|
config: EbookSearchConfig,
|
|
) -> tuple[dict[str, PhraseCandidate], int, int, int]:
|
|
"""Remove candidates that should not be persisted."""
|
|
min_raw_count = minimum_candidate_raw_count(config)
|
|
filtered: dict[str, PhraseCandidate] = {}
|
|
too_short = 0
|
|
too_rare = 0
|
|
too_common = 0
|
|
for phrase_norm, candidate in candidates.items():
|
|
if candidate.token_count < config.phrase_min_tokens:
|
|
too_short += 1
|
|
continue
|
|
if candidate.raw_count < min_raw_count:
|
|
too_rare += 1
|
|
continue
|
|
if is_most_common_word_phrase(phrase_norm):
|
|
too_common += 1
|
|
continue
|
|
filtered[phrase_norm] = candidate
|
|
return filtered, too_short, too_rare, too_common
|
|
|
|
|
|
def minimum_candidate_raw_count(config: EbookSearchConfig) -> int:
|
|
"""Return the minimum occurrence count required before storing a candidate."""
|
|
return max(config.phrase_raw_ngram_min_count, 1)
|
|
|
|
|
|
def is_most_common_word_phrase(phrase_norm: str) -> bool:
|
|
"""Return whether every token in a normalized phrase is a common word."""
|
|
tokens_ = phrase_norm.split()
|
|
common_words = get_most_common_words()
|
|
return bool(tokens_) and all(token in common_words for token in tokens_)
|
|
|
|
|
|
def score_candidate(candidate: PhraseCandidate, config: EbookSearchConfig) -> float:
|
|
"""Score a phrase candidate before LLM judging."""
|
|
score = source_score(candidate) + frequency_score(candidate, config) + token_count_score(candidate, config)
|
|
if candidate.phrase_norm in get_ignored_phrases():
|
|
score -= 100.0
|
|
if has_bad_start(candidate.phrase_norm):
|
|
score -= BAD_START_SCORE_PENALTY
|
|
if has_bad_end(candidate.phrase_norm):
|
|
score -= BAD_END_SCORE_PENALTY
|
|
return score
|
|
|
|
|
|
def has_bad_start(phrase_norm: str) -> bool:
|
|
"""Return whether a normalized phrase starts with a bad starting token."""
|
|
tokens_ = phrase_norm.split()
|
|
return bool(tokens_ and tokens_[0] in get_bad_starts())
|
|
|
|
|
|
def has_bad_end(phrase_norm: str) -> bool:
|
|
"""Return whether a normalized phrase ends with a bad ending token."""
|
|
tokens_ = phrase_norm.split()
|
|
return bool(tokens_ and tokens_[-1] in get_bad_ends())
|
|
|
|
|
|
def source_score(candidate: PhraseCandidate) -> float:
|
|
"""Return the score contribution from extraction sources."""
|
|
return sum(
|
|
weight
|
|
for enabled, weight in (
|
|
(candidate.source_yake, 2.0),
|
|
(candidate.source_spacy_ner, 2.5),
|
|
(candidate.source_spacy_noun_chunk, 1.5),
|
|
(candidate.source_capitalized, 2.0),
|
|
(candidate.source_raw_ngram, 0.5),
|
|
)
|
|
if enabled
|
|
)
|
|
|
|
|
|
def frequency_score(candidate: PhraseCandidate, config: EbookSearchConfig) -> float:
|
|
"""Return the score contribution from frequency and chapter spread."""
|
|
return sum(
|
|
weight
|
|
for count, threshold, weight in (
|
|
(candidate.raw_count, config.phrase_raw_count_score_threshold, 1.0),
|
|
(candidate.raw_count, config.phrase_raw_count_high_score_threshold, 1.0),
|
|
(candidate.chapter_count, config.phrase_chapter_count_score_threshold, 1.0),
|
|
(candidate.chapter_count, config.phrase_chapter_count_high_score_threshold, 1.0),
|
|
)
|
|
if count >= threshold
|
|
)
|
|
|
|
|
|
def token_count_score(candidate: PhraseCandidate, config: EbookSearchConfig) -> float:
|
|
"""Return the score contribution from phrase length."""
|
|
if candidate.token_count == 1:
|
|
return -0.5
|
|
if candidate.token_count in {2, 3, 4}:
|
|
return 0.5
|
|
if candidate.token_count > config.phrase_max_tokens:
|
|
return -1.0
|
|
return 0.0
|
|
|
|
|
|
def get_sample_contexts(book_text: str, phrase_norm: str, max_contexts: int = 5) -> list[str]:
|
|
"""Return normalized context snippets containing a candidate phrase."""
|
|
norm_text = normalize_text(book_text)
|
|
contexts: list[str] = []
|
|
start = 0
|
|
while len(contexts) < max_contexts:
|
|
index = norm_text.find(phrase_norm, start)
|
|
if index == -1:
|
|
break
|
|
left = max(0, index - 300)
|
|
right = min(len(norm_text), index + len(phrase_norm) + 300)
|
|
contexts.append(norm_text[left:right])
|
|
start = index + len(phrase_norm)
|
|
return contexts
|
|
|
|
|
|
def judge_candidate_with_llm(candidate: PhraseCandidate, config: EbookSearchConfig) -> LLMJudgment:
|
|
"""Ask the configured chat model to judge one pre-extracted candidate."""
|
|
payload = {
|
|
"phrase": candidate.phrase_norm,
|
|
"token_count": candidate.token_count,
|
|
"sources": candidate_source_names(candidate),
|
|
"raw_count": candidate.raw_count,
|
|
"chapter_count": candidate.chapter_count,
|
|
"contexts": list(candidate.sample_contexts),
|
|
}
|
|
messages = [
|
|
{
|
|
"role": "system",
|
|
"content": (
|
|
"Judge whether a candidate phrase from a book should be protected for RAG retrieval. "
|
|
"Do not extract new phrases. Reject common grammar fragments, ordinary nonspecific phrases, "
|
|
"unstable fragments, and phrases kept only because they are frequent. Keep people, places, "
|
|
"organizations, factions, events, technologies, fictional conditions, magic systems, formal titles, "
|
|
"named concepts, and recurring world-specific terms. Return only a JSON object with keys: keep, "
|
|
"canonical, category, aliases, confidence, importance, allow_nested, suppress_children, reason."
|
|
),
|
|
},
|
|
{"role": "user", "content": json.dumps(payload, ensure_ascii=True)},
|
|
]
|
|
return parse_llm_judgment(request_chat_completion(config, messages), config)
|
|
|
|
|
|
def candidate_source_names(candidate: PhraseCandidate) -> list[str]:
|
|
"""Return enabled source names for an extracted candidate."""
|
|
names: list[str] = []
|
|
if candidate.source_raw_ngram:
|
|
names.append("raw_ngram")
|
|
if candidate.source_yake:
|
|
names.append("yake")
|
|
if candidate.source_spacy_ner:
|
|
names.append("spacy_ner")
|
|
if candidate.source_spacy_noun_chunk:
|
|
names.append("spacy_noun_chunk")
|
|
if candidate.source_capitalized:
|
|
names.append("capitalized")
|
|
if candidate.source_metadata:
|
|
names.append("metadata")
|
|
return names
|
|
|
|
|
|
def parse_llm_judgment(content: str, config: EbookSearchConfig) -> LLMJudgment:
|
|
"""Parse and validate an LLM phrase-judge response."""
|
|
body = json.loads(extract_json_object(content))
|
|
if not isinstance(body, dict):
|
|
msg = "LLM phrase judge response is not a JSON object"
|
|
raise TypeError(msg)
|
|
|
|
aliases = body.get("aliases", ())
|
|
if not isinstance(aliases, list | tuple):
|
|
aliases = ()
|
|
return LLMJudgment(
|
|
keep=bool(body.get("keep", False)),
|
|
canonical=optional_text(body.get("canonical")),
|
|
category=optional_text(body.get("category")),
|
|
aliases=tuple(str(alias) for alias in aliases if isinstance(alias, str) and alias.strip()),
|
|
confidence=clamped_float(body.get("confidence"), default=0.0),
|
|
importance=clamped_float(body.get("importance"), default=0.5),
|
|
allow_nested=bool(body.get("allow_nested", config.phrase_default_allow_nested)),
|
|
suppress_children=bool(body.get("suppress_children", config.phrase_default_suppress_children)),
|
|
reason=optional_text(body.get("reason")),
|
|
)
|
|
|
|
|
|
def extract_json_object(content: str) -> str:
|
|
"""Extract a JSON object from plain or fenced model output."""
|
|
stripped = content.strip()
|
|
if stripped.startswith("{") and stripped.endswith("}"):
|
|
return stripped
|
|
match = JSON_OBJECT_RE.search(stripped)
|
|
if match is None:
|
|
msg = "LLM phrase judge response did not contain a JSON object"
|
|
raise ValueError(msg)
|
|
return match.group(0)
|
|
|
|
|
|
def optional_text(value: object) -> str | None:
|
|
"""Return stripped text for a nullable JSON value."""
|
|
if not isinstance(value, str):
|
|
return None
|
|
stripped = value.strip()
|
|
return stripped or None
|
|
|
|
|
|
def clamped_float(value: object, *, default: float) -> float:
|
|
"""Coerce a JSON number into the 0.0 to 1.0 range."""
|
|
if not isinstance(value, int | float):
|
|
return default
|
|
return min(max(float(value), 0.0), 1.0)
|
|
|
|
|
|
def extract_phrase_candidates_for_book(
|
|
book_text: str,
|
|
chapters: Sequence[str],
|
|
config: EbookSearchConfig,
|
|
*,
|
|
nlp: SpacyLanguage | None = None,
|
|
metadata: Mapping[str, object] | None = None,
|
|
) -> list[PhraseCandidate]:
|
|
"""Extract, score, and limit phrase candidates for one book."""
|
|
started_at = perf_counter()
|
|
logger.info(
|
|
"ebook_phrase_candidate_extract_start chapters=%s chars=%s min_tokens=%s max_tokens=%s max_candidates=%s",
|
|
len(chapters),
|
|
len(book_text),
|
|
config.phrase_min_tokens,
|
|
config.phrase_max_tokens,
|
|
config.protected_phrase_max_candidates_per_book,
|
|
)
|
|
raw_started_at = perf_counter()
|
|
raw = extract_raw_ngrams(book_text, config)
|
|
logger.info(
|
|
"ebook_phrase_candidate_extract_raw_complete candidates=%s duration_ms=%.1f",
|
|
len(raw),
|
|
(perf_counter() - raw_started_at) * 1000,
|
|
)
|
|
yake_started_at = perf_counter()
|
|
yake_candidates = extract_yake_candidates(book_text, config)
|
|
logger.info(
|
|
"ebook_phrase_candidate_extract_yake_complete candidates=%s duration_ms=%.1f",
|
|
len(yake_candidates),
|
|
(perf_counter() - yake_started_at) * 1000,
|
|
)
|
|
spacy_candidates: dict[str, PhraseCandidate] = {}
|
|
if nlp is not None:
|
|
spacy_started_at = perf_counter()
|
|
for chapter in chapters:
|
|
spacy_candidates = merge_candidate_sources(spacy_candidates, extract_spacy_candidates(chapter, nlp, config))
|
|
logger.info(
|
|
"ebook_phrase_candidate_extract_spacy_complete candidates=%s duration_ms=%.1f",
|
|
len(spacy_candidates),
|
|
(perf_counter() - spacy_started_at) * 1000,
|
|
)
|
|
capitalized_started_at = perf_counter()
|
|
capitalized = extract_capitalized_phrases(book_text, config)
|
|
logger.info(
|
|
"ebook_phrase_candidate_extract_capitalized_complete candidates=%s duration_ms=%.1f",
|
|
len(capitalized),
|
|
(perf_counter() - capitalized_started_at) * 1000,
|
|
)
|
|
metadata_candidates = extract_metadata_candidates(metadata, config)
|
|
|
|
candidates = merge_candidate_sources(raw, yake_candidates, spacy_candidates, capitalized, metadata_candidates)
|
|
enriched_started_at = perf_counter()
|
|
candidates = enrich_with_frequency_and_chapter_counts(candidates, chapters)
|
|
pre_filter_count = len(candidates)
|
|
candidates, filtered_too_short, filtered_too_rare, filtered_too_common = filter_storable_candidates(
|
|
candidates,
|
|
config,
|
|
)
|
|
for candidate in candidates.values():
|
|
candidate.candidate_score = score_candidate(candidate, config)
|
|
|
|
limited = sorted(candidates.values(), key=lambda item: item.candidate_score, reverse=True)[
|
|
: config.protected_phrase_max_candidates_per_book
|
|
]
|
|
logger.info(
|
|
"ebook_phrase_candidate_extract_complete raw=%s yake=%s spacy=%s capitalized=%s metadata=%s "
|
|
"merged=%s filtered_too_short=%s filtered_too_rare=%s filtered_too_common=%s min_uses=%s "
|
|
"storable=%s limited=%s enrich_score_ms=%.1f duration_ms=%.1f",
|
|
len(raw),
|
|
len(yake_candidates),
|
|
len(spacy_candidates),
|
|
len(capitalized),
|
|
len(metadata_candidates),
|
|
pre_filter_count,
|
|
filtered_too_short,
|
|
filtered_too_rare,
|
|
filtered_too_common,
|
|
minimum_candidate_raw_count(config),
|
|
len(candidates),
|
|
len(limited),
|
|
(perf_counter() - enriched_started_at) * 1000,
|
|
(perf_counter() - started_at) * 1000,
|
|
)
|
|
return limited
|
|
|
|
|
|
def generate_candidate_phrases_for_book(
|
|
session: Session,
|
|
book_id: int,
|
|
series_id: int | None,
|
|
book_text: str,
|
|
chapters: Sequence[str],
|
|
config: EbookSearchConfig,
|
|
*,
|
|
nlp: SpacyLanguage | None = None,
|
|
metadata: Mapping[str, object] | None = None,
|
|
) -> list[EbookCandidatePhrase]:
|
|
"""Extract and store candidate phrases for one book without LLM judging."""
|
|
started_at = perf_counter()
|
|
limited_candidates = extract_phrase_candidates_for_book(
|
|
book_text,
|
|
chapters,
|
|
config,
|
|
nlp=nlp,
|
|
metadata=metadata,
|
|
)
|
|
save_started_at = perf_counter()
|
|
pruned_count = prune_unstorable_unjudged_candidate_phrases(session, book_id, config)
|
|
logger.info(
|
|
"ebook_candidate_phrase_save_start book_id=%s candidates=%s pruned_unstorable=%s",
|
|
book_id,
|
|
len(limited_candidates),
|
|
pruned_count,
|
|
)
|
|
rows = [
|
|
save_candidate_to_db(session, book_id, series_id, candidate, judgment=None) for candidate in limited_candidates
|
|
]
|
|
session.flush()
|
|
logger.info(
|
|
"ebook_candidate_phrase_generation_complete book_id=%s candidates=%s save_ms=%.1f duration_ms=%.1f",
|
|
book_id,
|
|
len(rows),
|
|
(perf_counter() - save_started_at) * 1000,
|
|
(perf_counter() - started_at) * 1000,
|
|
)
|
|
return rows
|
|
|
|
|
|
def prune_unstorable_unjudged_candidate_phrases(
|
|
session: Session,
|
|
book_id: int,
|
|
config: EbookSearchConfig,
|
|
) -> int:
|
|
"""Delete old unjudged candidate rows that no longer satisfy storage filters."""
|
|
deleted = rowcount(
|
|
session.execute(
|
|
delete(EbookCandidatePhrase).where(
|
|
EbookCandidatePhrase.book_id == book_id,
|
|
EbookCandidatePhrase.llm_judged.is_(False),
|
|
or_(
|
|
EbookCandidatePhrase.token_count < config.phrase_min_tokens,
|
|
EbookCandidatePhrase.raw_count < minimum_candidate_raw_count(config),
|
|
),
|
|
)
|
|
)
|
|
)
|
|
if deleted:
|
|
logger.info(
|
|
"ebook_candidate_phrase_unstorable_pruned book_id=%s deleted=%s min_tokens=%s min_uses=%s",
|
|
book_id,
|
|
deleted,
|
|
config.phrase_min_tokens,
|
|
minimum_candidate_raw_count(config),
|
|
)
|
|
deleted += prune_common_word_candidate_phrases(session, book_id)
|
|
return deleted
|
|
|
|
|
|
def prune_common_word_candidate_phrases(session: Session, book_id: int) -> int:
|
|
"""Delete old unjudged candidates made only from common words."""
|
|
rows = list(
|
|
session.scalars(
|
|
select(EbookCandidatePhrase).where(
|
|
EbookCandidatePhrase.book_id == book_id,
|
|
EbookCandidatePhrase.llm_judged.is_(False),
|
|
)
|
|
)
|
|
)
|
|
ids = [row.id for row in rows if row.id is not None and is_most_common_word_phrase(row.phrase_norm)]
|
|
if not ids:
|
|
return 0
|
|
deleted = rowcount(session.execute(delete(EbookCandidatePhrase).where(EbookCandidatePhrase.id.in_(ids))))
|
|
logger.info("ebook_candidate_phrase_common_word_pruned book_id=%s deleted=%s", book_id, deleted)
|
|
return deleted
|
|
|
|
|
|
def build_protected_phrases_for_book(
|
|
session: Session,
|
|
book_id: int,
|
|
series_id: int | None,
|
|
book_text: str,
|
|
chapters: Sequence[str],
|
|
config: EbookSearchConfig,
|
|
*,
|
|
nlp: SpacyLanguage | None = None,
|
|
metadata: Mapping[str, object] | None = None,
|
|
) -> list[EbookProtectedPhrase]:
|
|
"""Extract candidates, judge stored candidates, and store accepted protected phrases."""
|
|
candidates = generate_candidate_phrases_for_book(
|
|
session,
|
|
book_id,
|
|
series_id,
|
|
book_text,
|
|
chapters,
|
|
config,
|
|
nlp=nlp,
|
|
metadata=metadata,
|
|
)
|
|
judged_count, protected = judge_candidate_phrases_for_book(
|
|
session,
|
|
book_id,
|
|
series_id,
|
|
book_text,
|
|
config,
|
|
)
|
|
|
|
logger.info(
|
|
"ebook_protected_phrase_build_complete book_id=%s candidates=%s judged=%s protected=%s",
|
|
book_id,
|
|
len(candidates),
|
|
judged_count,
|
|
len(protected),
|
|
)
|
|
return protected
|
|
|
|
|
|
def generate_candidate_phrases_for_books(
|
|
session: Session,
|
|
config: EbookSearchConfig,
|
|
*,
|
|
limit: int | None = None,
|
|
nlp: SpacyLanguage | None = None,
|
|
) -> PhraseCandidateGenerationResult:
|
|
"""Create or refresh candidate phrases for indexed books without calling the LLM judge."""
|
|
sources = list(session.scalars(select(EbookSource).order_by(EbookSource.id)))
|
|
books_seen = len(sources)
|
|
books_built = 0
|
|
candidate_count = 0
|
|
logger.info(
|
|
"ebook_candidate_phrase_generation_start books_seen=%s limit=%s min_tokens=%s max_tokens=%s "
|
|
"max_candidates_per_book=%s",
|
|
books_seen,
|
|
limit,
|
|
config.phrase_min_tokens,
|
|
config.phrase_max_tokens,
|
|
config.protected_phrase_max_candidates_per_book,
|
|
)
|
|
|
|
for book_number, source in enumerate(sources, start=1):
|
|
if limit is not None and books_built >= limit:
|
|
logger.info("ebook_candidate_phrase_generation_book_limit_reached limit=%s", limit)
|
|
break
|
|
book_started_at = perf_counter()
|
|
logger.info(
|
|
"ebook_candidate_phrase_generation_book_start source_id=%s book_number=%s books_seen=%s title=%r",
|
|
source.id,
|
|
book_number,
|
|
books_seen,
|
|
source.title,
|
|
)
|
|
|
|
try:
|
|
chapters = load_book_chapter_texts(session, source.id)
|
|
if not chapters:
|
|
logger.warning("ebook_candidate_phrase_generation_book_empty source_id=%s", source.id)
|
|
continue
|
|
book_text = "\n\n".join(chapters)
|
|
logger.info(
|
|
"ebook_candidate_phrase_generation_book_loaded source_id=%s chapters=%s chars=%s",
|
|
source.id,
|
|
len(chapters),
|
|
len(book_text),
|
|
)
|
|
|
|
candidates = generate_candidate_phrases_for_book(
|
|
session,
|
|
source.id,
|
|
series_id=None,
|
|
book_text=book_text,
|
|
chapters=chapters,
|
|
config=config,
|
|
nlp=nlp,
|
|
metadata=metadata_for_source(source),
|
|
)
|
|
session.commit()
|
|
except Exception:
|
|
session.rollback()
|
|
logger.exception(
|
|
"ebook_candidate_phrase_generation_book_failed source_id=%s book_number=%s books_seen=%s",
|
|
source.id,
|
|
book_number,
|
|
books_seen,
|
|
)
|
|
raise
|
|
books_built += 1
|
|
candidate_count += len(candidates)
|
|
logger.info(
|
|
"ebook_candidate_phrase_generation_book_committed source_id=%s candidates=%s books_built=%s "
|
|
"candidate_total=%s duration_ms=%.1f",
|
|
source.id,
|
|
len(candidates),
|
|
books_built,
|
|
candidate_count,
|
|
(perf_counter() - book_started_at) * 1000,
|
|
)
|
|
|
|
logger.info(
|
|
"ebook_candidate_phrase_generation_complete books_seen=%s books_built=%s candidate_total=%s",
|
|
books_seen,
|
|
books_built,
|
|
candidate_count,
|
|
)
|
|
return PhraseCandidateGenerationResult(
|
|
books_seen=books_seen,
|
|
books_built=books_built,
|
|
candidate_phrases=candidate_count,
|
|
)
|
|
|
|
|
|
def recalculate_candidate_phrases_for_book(
|
|
session: Session,
|
|
source: EbookSource,
|
|
config: EbookSearchConfig,
|
|
*,
|
|
nlp: SpacyLanguage | None = None,
|
|
) -> PhraseRecalculationResult:
|
|
"""Remove all book phrase data, regenerate candidates, and commit the completed book."""
|
|
started_at = perf_counter()
|
|
logger.info(
|
|
"ebook_candidate_phrase_recalculation_start source_id=%s title=%r",
|
|
source.id,
|
|
source.title,
|
|
)
|
|
try:
|
|
deleted = delete_phrase_data_for_book(session, source.id)
|
|
chapters = load_book_chapter_texts(session, source.id)
|
|
if not chapters:
|
|
logger.warning("ebook_candidate_phrase_recalculation_book_empty source_id=%s", source.id)
|
|
session.commit()
|
|
return PhraseRecalculationResult(
|
|
book_id=source.id,
|
|
deleted_candidates=deleted.deleted_candidates,
|
|
deleted_protected_phrases=deleted.deleted_protected_phrases,
|
|
deleted_aliases=deleted.deleted_aliases,
|
|
deleted_mentions=deleted.deleted_mentions,
|
|
candidate_phrases=0,
|
|
)
|
|
|
|
candidates = generate_candidate_phrases_for_book(
|
|
session,
|
|
source.id,
|
|
series_id=None,
|
|
book_text="\n\n".join(chapters),
|
|
chapters=chapters,
|
|
config=config,
|
|
nlp=nlp,
|
|
metadata=metadata_for_source(source),
|
|
)
|
|
session.commit()
|
|
except Exception:
|
|
session.rollback()
|
|
logger.exception("ebook_candidate_phrase_recalculation_failed source_id=%s", source.id)
|
|
raise
|
|
|
|
result = PhraseRecalculationResult(
|
|
book_id=source.id,
|
|
deleted_candidates=deleted.deleted_candidates,
|
|
deleted_protected_phrases=deleted.deleted_protected_phrases,
|
|
deleted_aliases=deleted.deleted_aliases,
|
|
deleted_mentions=deleted.deleted_mentions,
|
|
candidate_phrases=len(candidates),
|
|
)
|
|
logger.info(
|
|
"ebook_candidate_phrase_recalculation_complete source_id=%s deleted_candidates=%s "
|
|
"deleted_protected=%s deleted_aliases=%s deleted_mentions=%s candidates=%s duration_ms=%.1f",
|
|
source.id,
|
|
result.deleted_candidates,
|
|
result.deleted_protected_phrases,
|
|
result.deleted_aliases,
|
|
result.deleted_mentions,
|
|
result.candidate_phrases,
|
|
(perf_counter() - started_at) * 1000,
|
|
)
|
|
return result
|
|
|
|
|
|
def delete_phrase_data_for_book(session: Session, book_id: int) -> PhraseRecalculationResult:
|
|
"""Delete all candidate, protected, alias, and mention phrase data for one book."""
|
|
protected_ids = list(
|
|
session.scalars(select(EbookProtectedPhrase.id).where(EbookProtectedPhrase.book_id == book_id))
|
|
)
|
|
deleted_aliases = 0
|
|
if protected_ids:
|
|
deleted_aliases = rowcount(
|
|
session.execute(delete(EbookPhraseAlias).where(EbookPhraseAlias.phrase_id.in_(protected_ids)))
|
|
)
|
|
|
|
deleted_mentions = rowcount(
|
|
session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.book_id == book_id))
|
|
)
|
|
if protected_ids:
|
|
deleted_mentions += rowcount(
|
|
session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.phrase_id.in_(protected_ids)))
|
|
)
|
|
|
|
deleted_protected = rowcount(
|
|
session.execute(delete(EbookProtectedPhrase).where(EbookProtectedPhrase.book_id == book_id))
|
|
)
|
|
deleted_candidates = rowcount(
|
|
session.execute(delete(EbookCandidatePhrase).where(EbookCandidatePhrase.book_id == book_id))
|
|
)
|
|
session.flush()
|
|
logger.info(
|
|
"ebook_candidate_phrase_data_deleted book_id=%s candidates=%s protected=%s aliases=%s mentions=%s",
|
|
book_id,
|
|
deleted_candidates,
|
|
deleted_protected,
|
|
deleted_aliases,
|
|
deleted_mentions,
|
|
)
|
|
return PhraseRecalculationResult(
|
|
book_id=book_id,
|
|
deleted_candidates=deleted_candidates,
|
|
deleted_protected_phrases=deleted_protected,
|
|
deleted_aliases=deleted_aliases,
|
|
deleted_mentions=deleted_mentions,
|
|
candidate_phrases=0,
|
|
)
|
|
|
|
|
|
def rowcount(result: object) -> int:
|
|
"""Return a safe integer rowcount from a SQLAlchemy execution result."""
|
|
count = getattr(result, "rowcount", 0)
|
|
return int(count if count is not None and count >= 0 else 0)
|
|
|
|
|
|
def judge_candidate_phrases_for_books(
|
|
session: Session,
|
|
config: EbookSearchConfig,
|
|
*,
|
|
limit: int | None = None,
|
|
) -> PhraseJudgmentBackfillResult:
|
|
"""Judge stored candidate phrases and promote accepted phrases for indexed books."""
|
|
sources = list(session.scalars(select(EbookSource).order_by(EbookSource.id)))
|
|
books_seen = len(sources)
|
|
books_judged = 0
|
|
candidates_judged = 0
|
|
protected_count = 0
|
|
mention_count = 0
|
|
logger.info(
|
|
"ebook_candidate_phrase_judgment_start books_seen=%s limit=%s llm_candidates_per_book=%s "
|
|
"target_protected_per_book=%s confidence_threshold=%.2f min_tokens=%s min_uses=%s",
|
|
books_seen,
|
|
limit,
|
|
config.protected_phrase_llm_candidates_per_book,
|
|
config.phrase_target_protected_per_book,
|
|
config.protected_phrase_confidence_threshold,
|
|
config.phrase_min_tokens,
|
|
minimum_candidate_raw_count(config),
|
|
)
|
|
|
|
for book_number, source in enumerate(sources, start=1):
|
|
if limit is not None and books_judged >= limit:
|
|
logger.info("ebook_candidate_phrase_judgment_book_limit_reached limit=%s", limit)
|
|
break
|
|
book_started_at = perf_counter()
|
|
pruned_count = prune_common_word_candidate_phrases(session, source.id)
|
|
if pruned_count:
|
|
session.flush()
|
|
|
|
unjudged_count = session.scalar(
|
|
select(func.count(EbookCandidatePhrase.id)).where(
|
|
EbookCandidatePhrase.book_id == source.id,
|
|
EbookCandidatePhrase.llm_judged.is_(False),
|
|
EbookCandidatePhrase.token_count >= config.phrase_min_tokens,
|
|
EbookCandidatePhrase.raw_count >= minimum_candidate_raw_count(config),
|
|
)
|
|
)
|
|
if not unjudged_count:
|
|
logger.info(
|
|
"ebook_candidate_phrase_judgment_book_skip_no_unjudged source_id=%s book_number=%s books_seen=%s",
|
|
source.id,
|
|
book_number,
|
|
books_seen,
|
|
)
|
|
continue
|
|
logger.info(
|
|
"ebook_candidate_phrase_judgment_book_start source_id=%s book_number=%s books_seen=%s "
|
|
"title=%r unjudged=%s",
|
|
source.id,
|
|
book_number,
|
|
books_seen,
|
|
source.title,
|
|
unjudged_count,
|
|
)
|
|
|
|
try:
|
|
chapters = load_book_chapter_texts(session, source.id)
|
|
if not chapters:
|
|
logger.warning("ebook_candidate_phrase_judgment_book_empty source_id=%s", source.id)
|
|
continue
|
|
|
|
judged, protected = judge_candidate_phrases_for_book(
|
|
session,
|
|
source.id,
|
|
series_id=None,
|
|
book_text="\n\n".join(chapters),
|
|
config=config,
|
|
)
|
|
if judged == 0:
|
|
logger.info(
|
|
"ebook_candidate_phrase_judgment_book_skip_no_judgments source_id=%s unjudged=%s",
|
|
source.id,
|
|
unjudged_count,
|
|
)
|
|
continue
|
|
|
|
mentions = index_chunk_phrase_mentions_for_book(session, source.id, config) if protected else 0
|
|
session.commit()
|
|
except Exception:
|
|
session.rollback()
|
|
logger.exception(
|
|
"ebook_candidate_phrase_judgment_book_failed source_id=%s book_number=%s books_seen=%s",
|
|
source.id,
|
|
book_number,
|
|
books_seen,
|
|
)
|
|
raise
|
|
books_judged += 1
|
|
candidates_judged += judged
|
|
protected_count += len(protected)
|
|
mention_count += mentions
|
|
logger.info(
|
|
"ebook_candidate_phrase_judgment_book_committed source_id=%s judged=%s protected=%s mentions=%s "
|
|
"books_judged=%s candidates_judged_total=%s protected_total=%s duration_ms=%.1f",
|
|
source.id,
|
|
judged,
|
|
len(protected),
|
|
mentions,
|
|
books_judged,
|
|
candidates_judged,
|
|
protected_count,
|
|
(perf_counter() - book_started_at) * 1000,
|
|
)
|
|
|
|
logger.info(
|
|
"ebook_candidate_phrase_judgment_complete books_seen=%s books_judged=%s candidates_judged=%s "
|
|
"protected=%s mentions=%s",
|
|
books_seen,
|
|
books_judged,
|
|
candidates_judged,
|
|
protected_count,
|
|
mention_count,
|
|
)
|
|
return PhraseJudgmentBackfillResult(
|
|
books_seen=books_seen,
|
|
books_judged=books_judged,
|
|
candidates_judged=candidates_judged,
|
|
protected_phrases=protected_count,
|
|
phrase_mentions=mention_count,
|
|
)
|
|
|
|
|
|
def judge_candidate_phrases_for_book(
|
|
session: Session,
|
|
book_id: int,
|
|
series_id: int | None,
|
|
book_text: str,
|
|
config: EbookSearchConfig,
|
|
) -> tuple[int, list[EbookProtectedPhrase]]:
|
|
"""Judge unjudged candidate phrase rows for one book."""
|
|
judgment_limit = config.protected_phrase_llm_candidates_per_book
|
|
if judgment_limit <= 0:
|
|
logger.info("ebook_candidate_phrase_judgment_skipped_llm_limit_zero book_id=%s", book_id)
|
|
return 0, []
|
|
|
|
existing_protected = session.scalar(
|
|
select(func.count(EbookProtectedPhrase.id)).where(EbookProtectedPhrase.book_id == book_id)
|
|
)
|
|
target_remaining: int | None = None
|
|
if config.phrase_target_protected_per_book > 0:
|
|
target_remaining = max(config.phrase_target_protected_per_book - int(existing_protected or 0), 0)
|
|
if target_remaining == 0:
|
|
logger.info(
|
|
"ebook_candidate_phrase_judgment_skipped_target_met book_id=%s existing_protected=%s target=%s",
|
|
book_id,
|
|
existing_protected,
|
|
config.phrase_target_protected_per_book,
|
|
)
|
|
return 0, []
|
|
|
|
rows = list(
|
|
session.scalars(
|
|
select(EbookCandidatePhrase)
|
|
.where(
|
|
EbookCandidatePhrase.book_id == book_id,
|
|
EbookCandidatePhrase.llm_judged.is_(False),
|
|
EbookCandidatePhrase.token_count >= config.phrase_min_tokens,
|
|
EbookCandidatePhrase.raw_count >= minimum_candidate_raw_count(config),
|
|
)
|
|
.order_by(
|
|
EbookCandidatePhrase.candidate_score.desc(),
|
|
EbookCandidatePhrase.raw_count.desc(),
|
|
EbookCandidatePhrase.id,
|
|
)
|
|
.limit(judgment_limit)
|
|
)
|
|
)
|
|
rows = [row for row in rows if not is_most_common_word_phrase(row.phrase_norm)]
|
|
|
|
logger.info(
|
|
"ebook_candidate_phrase_judgment_candidates_loaded book_id=%s candidates=%s existing_protected=%s "
|
|
"target_remaining=%s judgment_limit=%s",
|
|
book_id,
|
|
len(rows),
|
|
existing_protected,
|
|
target_remaining,
|
|
judgment_limit,
|
|
)
|
|
judged_count = 0
|
|
protected: list[EbookProtectedPhrase] = []
|
|
for row_number, row in enumerate(rows, start=1):
|
|
row_started_at = perf_counter()
|
|
candidate = phrase_candidate_from_row(row)
|
|
candidate.sample_contexts = tuple(row.sample_contexts or get_sample_contexts(book_text, candidate.phrase_norm))
|
|
logger.info(
|
|
"ebook_candidate_phrase_judgment_candidate_start book_id=%s candidate_id=%s row_number=%s rows=%s "
|
|
"phrase=%r score=%.3f raw_count=%s chapter_count=%s",
|
|
book_id,
|
|
row.id,
|
|
row_number,
|
|
len(rows),
|
|
candidate.phrase_norm,
|
|
candidate.candidate_score,
|
|
candidate.raw_count,
|
|
candidate.chapter_count,
|
|
)
|
|
judgment = judge_candidate_with_llm(candidate, config)
|
|
candidate_row = save_candidate_to_db(session, book_id, series_id, candidate, judgment=judgment)
|
|
judged_count += 1
|
|
logger.info(
|
|
"ebook_candidate_phrase_judgment_candidate_complete book_id=%s candidate_id=%s phrase=%r keep=%s "
|
|
"confidence=%.3f importance=%.3f category=%r duration_ms=%.1f",
|
|
book_id,
|
|
row.id,
|
|
candidate.phrase_norm,
|
|
judgment.keep,
|
|
judgment.confidence,
|
|
judgment.importance,
|
|
judgment.category,
|
|
(perf_counter() - row_started_at) * 1000,
|
|
)
|
|
if not judgment.keep or judgment.confidence < config.protected_phrase_confidence_threshold:
|
|
continue
|
|
accepted_norm = normalize_text(judgment.canonical or candidate.phrase_text)
|
|
accepted_token_count = len(accepted_norm.split())
|
|
if accepted_token_count < config.phrase_min_tokens:
|
|
logger.info(
|
|
"ebook_candidate_phrase_judgment_candidate_skip_short_canonical book_id=%s candidate_id=%s "
|
|
"phrase=%r canonical=%r token_count=%s min_tokens=%s",
|
|
book_id,
|
|
row.id,
|
|
candidate.phrase_norm,
|
|
accepted_norm,
|
|
accepted_token_count,
|
|
config.phrase_min_tokens,
|
|
)
|
|
continue
|
|
if is_most_common_word_phrase(accepted_norm):
|
|
logger.info(
|
|
"ebook_candidate_phrase_judgment_candidate_skip_common_canonical book_id=%s candidate_id=%s "
|
|
"phrase=%r canonical=%r",
|
|
book_id,
|
|
row.id,
|
|
candidate.phrase_norm,
|
|
accepted_norm,
|
|
)
|
|
continue
|
|
protected.append(upsert_protected_phrase(session, book_id, series_id, candidate, judgment, candidate_row))
|
|
if target_remaining is not None and len(protected) >= target_remaining:
|
|
break
|
|
|
|
session.flush()
|
|
return judged_count, protected
|
|
|
|
|
|
def phrase_candidate_from_row(row: EbookCandidatePhrase) -> PhraseCandidate:
|
|
"""Recreate an in-memory candidate from a persisted candidate row."""
|
|
return PhraseCandidate(
|
|
phrase_text=row.phrase_text,
|
|
phrase_norm=row.phrase_norm,
|
|
token_count=row.token_count,
|
|
source_raw_ngram=row.source_raw_ngram,
|
|
source_yake=row.source_yake,
|
|
source_spacy_ner=row.source_spacy_ner,
|
|
source_spacy_noun_chunk=row.source_spacy_noun_chunk,
|
|
source_capitalized=row.source_capitalized,
|
|
source_metadata=row.source_metadata,
|
|
spacy_label=row.spacy_label,
|
|
raw_count=row.raw_count,
|
|
chapter_count=row.chapter_count,
|
|
yake_score=row.yake_score,
|
|
candidate_score=row.candidate_score,
|
|
sample_contexts=tuple(row.sample_contexts or ()),
|
|
)
|
|
|
|
|
|
def metadata_for_source(source: EbookSource) -> dict[str, object | None]:
|
|
"""Return phrase extraction metadata for one indexed source."""
|
|
return {
|
|
"title": source.title,
|
|
"author": source.author,
|
|
"language": source.language,
|
|
"publisher": source.publisher,
|
|
"identifier": source.identifier,
|
|
}
|
|
|
|
|
|
def build_missing_protected_phrases(
|
|
session: Session,
|
|
config: EbookSearchConfig,
|
|
*,
|
|
limit: int | None = None,
|
|
nlp: SpacyLanguage | None = None,
|
|
) -> ProtectedPhraseBackfillResult:
|
|
"""Create protected phrases and phrase mentions for indexed books that do not have them yet."""
|
|
sources = list(session.scalars(select(EbookSource).order_by(EbookSource.id)))
|
|
books_seen = len(sources)
|
|
books_built = 0
|
|
protected_count = 0
|
|
mention_count = 0
|
|
logger.info(
|
|
"ebook_protected_phrase_backfill_start books_seen=%s limit=%s",
|
|
books_seen,
|
|
limit,
|
|
)
|
|
|
|
for book_number, source in enumerate(sources, start=1):
|
|
if limit is not None and books_built >= limit:
|
|
logger.info("ebook_protected_phrase_backfill_book_limit_reached limit=%s", limit)
|
|
break
|
|
book_started_at = perf_counter()
|
|
existing_count = session.scalar(
|
|
select(func.count(EbookProtectedPhrase.id)).where(EbookProtectedPhrase.book_id == source.id)
|
|
)
|
|
if existing_count:
|
|
logger.info(
|
|
"ebook_protected_phrase_backfill_book_skip_existing source_id=%s book_number=%s books_seen=%s "
|
|
"existing=%s",
|
|
source.id,
|
|
book_number,
|
|
books_seen,
|
|
existing_count,
|
|
)
|
|
continue
|
|
logger.info(
|
|
"ebook_protected_phrase_backfill_book_start source_id=%s book_number=%s books_seen=%s title=%r",
|
|
source.id,
|
|
book_number,
|
|
books_seen,
|
|
source.title,
|
|
)
|
|
|
|
try:
|
|
chapters = load_book_chapter_texts(session, source.id)
|
|
if not chapters:
|
|
logger.warning("ebook_protected_phrase_backfill_book_empty source_id=%s", source.id)
|
|
continue
|
|
|
|
protected = build_protected_phrases_for_book(
|
|
session,
|
|
source.id,
|
|
series_id=None,
|
|
book_text="\n\n".join(chapters),
|
|
chapters=chapters,
|
|
config=config,
|
|
nlp=nlp,
|
|
metadata=metadata_for_source(source),
|
|
)
|
|
mentions = index_chunk_phrase_mentions_for_book(session, source.id, config)
|
|
session.commit()
|
|
except Exception:
|
|
session.rollback()
|
|
logger.exception(
|
|
"ebook_protected_phrase_backfill_book_failed source_id=%s book_number=%s books_seen=%s",
|
|
source.id,
|
|
book_number,
|
|
books_seen,
|
|
)
|
|
raise
|
|
books_built += 1
|
|
protected_count += len(protected)
|
|
mention_count += mentions
|
|
logger.info(
|
|
"ebook_protected_phrase_backfill_book_committed source_id=%s protected=%s mentions=%s books_built=%s "
|
|
"protected_total=%s mention_total=%s duration_ms=%.1f",
|
|
source.id,
|
|
len(protected),
|
|
mentions,
|
|
books_built,
|
|
protected_count,
|
|
mention_count,
|
|
(perf_counter() - book_started_at) * 1000,
|
|
)
|
|
|
|
logger.info(
|
|
"ebook_protected_phrase_backfill_complete books_seen=%s books_built=%s protected=%s mentions=%s",
|
|
books_seen,
|
|
books_built,
|
|
protected_count,
|
|
mention_count,
|
|
)
|
|
return ProtectedPhraseBackfillResult(
|
|
books_seen=books_seen,
|
|
books_built=books_built,
|
|
protected_phrases=protected_count,
|
|
phrase_mentions=mention_count,
|
|
)
|
|
|
|
|
|
def load_book_chapter_texts(session: Session, book_id: int) -> list[str]:
|
|
"""Reconstruct chapter-like text blocks from indexed chunks for phrase extraction."""
|
|
rows = session.execute(
|
|
select(EbookChunk.chapter_id, EbookChunk.text)
|
|
.where(EbookChunk.source_id == book_id)
|
|
.order_by(EbookChunk.chunk_index)
|
|
)
|
|
chapters: list[str] = []
|
|
current_chapter_id: int | None = None
|
|
current_parts: list[str] = []
|
|
have_current = False
|
|
|
|
for chapter_id, text in rows:
|
|
if have_current and chapter_id != current_chapter_id:
|
|
chapter_text = "\n\n".join(current_parts).strip()
|
|
if chapter_text:
|
|
chapters.append(chapter_text)
|
|
current_parts = []
|
|
current_chapter_id = chapter_id
|
|
current_parts.append(str(text))
|
|
have_current = True
|
|
|
|
if current_parts:
|
|
chapter_text = "\n\n".join(current_parts).strip()
|
|
if chapter_text:
|
|
chapters.append(chapter_text)
|
|
return chapters
|
|
|
|
|
|
def save_candidate_to_db(
|
|
session: Session,
|
|
book_id: int,
|
|
series_id: int | None,
|
|
candidate: PhraseCandidate,
|
|
*,
|
|
judgment: LLMJudgment | None,
|
|
) -> EbookCandidatePhrase:
|
|
"""Insert or update one candidate phrase row."""
|
|
row = session.scalar(
|
|
select(EbookCandidatePhrase).where(
|
|
EbookCandidatePhrase.book_id == book_id,
|
|
EbookCandidatePhrase.phrase_norm == candidate.phrase_norm,
|
|
)
|
|
)
|
|
if row is None:
|
|
row = EbookCandidatePhrase(book_id=book_id, phrase_norm=candidate.phrase_norm)
|
|
row.llm_judged = False
|
|
session.add(row)
|
|
row.series_id = series_id
|
|
row.phrase_text = candidate.phrase_text
|
|
row.token_count = candidate.token_count
|
|
row.source_raw_ngram = candidate.source_raw_ngram
|
|
row.source_yake = candidate.source_yake
|
|
row.source_spacy_ner = candidate.source_spacy_ner
|
|
row.source_spacy_noun_chunk = candidate.source_spacy_noun_chunk
|
|
row.source_capitalized = candidate.source_capitalized
|
|
row.source_metadata = candidate.source_metadata
|
|
row.spacy_label = candidate.spacy_label
|
|
row.raw_count = candidate.raw_count
|
|
row.chapter_count = candidate.chapter_count
|
|
row.yake_score = candidate.yake_score
|
|
row.candidate_score = candidate.candidate_score
|
|
if candidate.sample_contexts:
|
|
row.sample_contexts = list(candidate.sample_contexts)
|
|
if judgment is not None:
|
|
row.llm_judged = True
|
|
row.llm_keep = judgment.keep
|
|
row.llm_confidence = judgment.confidence
|
|
row.llm_category = judgment.category
|
|
row.llm_reason = judgment.reason
|
|
return row
|
|
|
|
|
|
def upsert_protected_phrase(
|
|
session: Session,
|
|
book_id: int,
|
|
series_id: int | None,
|
|
candidate: PhraseCandidate,
|
|
judgment: LLMJudgment,
|
|
source_candidate: EbookCandidatePhrase,
|
|
) -> EbookProtectedPhrase:
|
|
"""Insert or update one accepted protected phrase and its aliases."""
|
|
phrase_text = judgment.canonical or candidate.phrase_text
|
|
phrase_norm = normalize_text(phrase_text)
|
|
if not phrase_norm:
|
|
msg = f"Protected phrase normalized to empty text: {phrase_text!r}"
|
|
raise ValueError(msg)
|
|
|
|
row = session.scalar(
|
|
select(EbookProtectedPhrase).where(
|
|
EbookProtectedPhrase.book_id == book_id,
|
|
EbookProtectedPhrase.phrase_norm == phrase_norm,
|
|
)
|
|
)
|
|
if row is None:
|
|
row = EbookProtectedPhrase(book_id=book_id, phrase_norm=phrase_norm)
|
|
session.add(row)
|
|
row.series_id = series_id
|
|
row.phrase_text = phrase_text
|
|
row.canonical_id = make_canonical_id(judgment, phrase_norm)
|
|
row.phrase_type = judgment.category
|
|
row.token_count = len(phrase_norm.split())
|
|
row.confidence = judgment.confidence
|
|
row.importance = judgment.importance
|
|
row.allow_nested = judgment.allow_nested
|
|
row.suppress_children = judgment.suppress_children
|
|
row.source_candidate_id = source_candidate.id
|
|
session.flush()
|
|
|
|
for alias_text in judgment.aliases:
|
|
upsert_phrase_alias(session, row, alias_text)
|
|
return row
|
|
|
|
|
|
def upsert_phrase_alias(session: Session, phrase: EbookProtectedPhrase, alias_text: str) -> EbookPhraseAlias | None:
|
|
"""Insert or update one protected phrase alias."""
|
|
alias_norm = normalize_text(alias_text)
|
|
if not alias_norm or alias_norm == phrase.phrase_norm:
|
|
return None
|
|
|
|
row = session.scalar(
|
|
select(EbookPhraseAlias).where(
|
|
EbookPhraseAlias.phrase_id == phrase.id,
|
|
EbookPhraseAlias.alias_norm == alias_norm,
|
|
)
|
|
)
|
|
if row is None:
|
|
row = EbookPhraseAlias(phrase_id=phrase.id, alias_norm=alias_norm)
|
|
session.add(row)
|
|
row.alias_text = alias_text
|
|
row.confidence = 1.0
|
|
return row
|
|
|
|
|
|
def make_canonical_id(judgment: LLMJudgment, phrase_norm: str) -> str:
|
|
"""Create a deterministic canonical id from a judgment category and phrase."""
|
|
category = slugify_identifier(judgment.category or "phrase")
|
|
phrase_slug = slugify_identifier(phrase_norm)
|
|
return f"{category}:{phrase_slug}"
|
|
|
|
|
|
def slugify_identifier(value: str) -> str:
|
|
"""Normalize text for use inside a canonical id."""
|
|
slug = re.sub(r"[^a-z0-9]+", "_", normalize_text(value).replace("'", ""))
|
|
return slug.strip("_") or "unknown"
|
|
|
|
|
|
def load_phrase_lookup(
|
|
session: Session,
|
|
config: EbookSearchConfig,
|
|
*,
|
|
book_id: int | None = None,
|
|
series_id: int | None = None,
|
|
) -> PhraseLookup:
|
|
"""Load protected phrases and aliases into RAM lookup maps."""
|
|
norm_to_ids: defaultdict[str, list[int]] = defaultdict(list)
|
|
alias_to_ids: defaultdict[str, list[int]] = defaultdict(list)
|
|
max_tokens = config.phrase_max_tokens
|
|
|
|
phrase_statement = select(
|
|
EbookProtectedPhrase.id,
|
|
EbookProtectedPhrase.phrase_norm,
|
|
EbookProtectedPhrase.token_count,
|
|
)
|
|
scope_filter = protected_phrase_scope_filter(book_id=book_id, series_id=series_id)
|
|
if scope_filter is not None:
|
|
phrase_statement = phrase_statement.where(scope_filter)
|
|
|
|
for row in session.execute(phrase_statement):
|
|
phrase_id = int(row.id)
|
|
phrase_norm = str(row.phrase_norm)
|
|
norm_to_ids[phrase_norm].append(phrase_id)
|
|
max_tokens = max(max_tokens, int(row.token_count))
|
|
|
|
alias_statement = select(
|
|
EbookPhraseAlias.alias_norm,
|
|
EbookPhraseAlias.phrase_id,
|
|
).join(EbookProtectedPhrase, EbookProtectedPhrase.id == EbookPhraseAlias.phrase_id)
|
|
if scope_filter is not None:
|
|
alias_statement = alias_statement.where(scope_filter)
|
|
|
|
for row in session.execute(alias_statement):
|
|
alias_norm = str(row.alias_norm)
|
|
alias_to_ids[alias_norm].append(int(row.phrase_id))
|
|
max_tokens = max(max_tokens, len(alias_norm.split()))
|
|
|
|
return PhraseLookup(
|
|
norm_to_phrase_ids={key: tuple(values) for key, values in norm_to_ids.items()},
|
|
alias_to_phrase_ids={key: tuple(values) for key, values in alias_to_ids.items()},
|
|
min_tokens=config.phrase_min_tokens,
|
|
max_tokens=max_tokens,
|
|
)
|
|
|
|
|
|
def protected_phrase_scope_filter(*, book_id: int | None, series_id: int | None) -> object | None:
|
|
"""Build a SQLAlchemy filter for optional phrase book and series scope."""
|
|
conditions = []
|
|
if book_id is not None:
|
|
conditions.append(or_(EbookProtectedPhrase.book_id.is_(None), EbookProtectedPhrase.book_id == book_id))
|
|
if series_id is not None:
|
|
conditions.append(or_(EbookProtectedPhrase.series_id.is_(None), EbookProtectedPhrase.series_id == series_id))
|
|
if not conditions:
|
|
return None
|
|
return and_(*conditions)
|
|
|
|
|
|
def generate_query_ngrams(
|
|
tokens_: Sequence[str],
|
|
min_n: int,
|
|
max_n: int,
|
|
) -> Iterator[tuple[str, int, int]]:
|
|
"""Generate normalized query windows from longest to shortest."""
|
|
capped_max_n = min(max_n, len(tokens_))
|
|
for ngram_size in range(capped_max_n, min_n - 1, -1):
|
|
for start in range(len(tokens_) - ngram_size + 1):
|
|
end = start + ngram_size
|
|
phrase_norm = " ".join(tokens_[start:end])
|
|
if phrase_norm in get_ignored_phrases():
|
|
continue
|
|
yield phrase_norm, start, end
|
|
|
|
|
|
def detect_phrase_candidates(query_text: str, lookup: PhraseLookup) -> list[PhraseMatch]:
|
|
"""Detect protected phrase windows in a user query using RAM hash lookups."""
|
|
return detect_phrase_candidates_from_tokens(tokenize_with_offsets(query_text), lookup)
|
|
|
|
|
|
def detect_phrase_candidates_in_text(text: str, lookup: PhraseLookup) -> list[PhraseMatch]:
|
|
"""Detect protected phrase windows in arbitrary text with character offsets."""
|
|
return detect_phrase_candidates_from_tokens(tokenize_with_offsets(text), lookup)
|
|
|
|
|
|
def detect_phrase_candidates_from_tokens(tokens_: Sequence[NormalizedToken], lookup: PhraseLookup) -> list[PhraseMatch]:
|
|
"""Detect protected phrase windows from already-normalized tokens."""
|
|
matches: list[PhraseMatch] = []
|
|
seen: set[tuple[int | None, str, int, int]] = set()
|
|
token_texts = [token.text for token in tokens_]
|
|
for phrase_norm, start, end in generate_query_ngrams(token_texts, min_n=lookup.min_tokens, max_n=lookup.max_tokens):
|
|
phrase_ids = lookup.norm_to_phrase_ids.get(phrase_norm, ())
|
|
alias_ids = lookup.alias_to_phrase_ids.get(phrase_norm, ())
|
|
for phrase_id in (*phrase_ids, *alias_ids):
|
|
key = (phrase_id, phrase_norm, start, end)
|
|
if key in seen:
|
|
continue
|
|
seen.add(key)
|
|
matches.append(
|
|
PhraseMatch(
|
|
phrase_norm=phrase_norm,
|
|
phrase_id=phrase_id,
|
|
start_token=start,
|
|
end_token=end,
|
|
token_count=end - start,
|
|
start_char=tokens_[start].start_char,
|
|
end_char=tokens_[end - 1].end_char,
|
|
)
|
|
)
|
|
return matches
|
|
|
|
|
|
def hydrate_matches(session: Session, matches: Sequence[PhraseMatch]) -> list[HydratedPhraseMatch]:
|
|
"""Fetch protected phrase metadata for raw phrase matches."""
|
|
if not matches:
|
|
return []
|
|
|
|
phrase_ids = sorted({match.phrase_id for match in matches if match.phrase_id is not None})
|
|
if not phrase_ids:
|
|
return []
|
|
|
|
rows = {
|
|
row.id: row
|
|
for row in session.scalars(select(EbookProtectedPhrase).where(EbookProtectedPhrase.id.in_(phrase_ids)))
|
|
}
|
|
hydrated: list[HydratedPhraseMatch] = []
|
|
for match in matches:
|
|
if match.phrase_id is None:
|
|
continue
|
|
phrase = rows.get(match.phrase_id)
|
|
if phrase is None:
|
|
continue
|
|
hydrated.append(
|
|
HydratedPhraseMatch(
|
|
phrase_id=phrase.id,
|
|
matched_norm=match.phrase_norm,
|
|
phrase_text=phrase.phrase_text,
|
|
phrase_norm=phrase.phrase_norm,
|
|
canonical_id=phrase.canonical_id,
|
|
phrase_type=phrase.phrase_type,
|
|
token_count=match.token_count,
|
|
confidence=phrase.confidence,
|
|
importance=phrase.importance,
|
|
allow_nested=phrase.allow_nested,
|
|
suppress_children=phrase.suppress_children,
|
|
start_token=match.start_token,
|
|
end_token=match.end_token,
|
|
start_char=match.start_char,
|
|
end_char=match.end_char,
|
|
book_id=phrase.book_id,
|
|
series_id=phrase.series_id,
|
|
)
|
|
)
|
|
return hydrated
|
|
|
|
|
|
def overlaps(first: HydratedPhraseMatch, second: HydratedPhraseMatch) -> bool:
|
|
"""Return whether two token spans overlap."""
|
|
return not (first.end_token <= second.start_token or first.start_token >= second.end_token)
|
|
|
|
|
|
def is_inside(child: HydratedPhraseMatch, parent: HydratedPhraseMatch) -> bool:
|
|
"""Return whether one token span is strictly inside another."""
|
|
return (
|
|
child.start_token >= parent.start_token
|
|
and child.end_token <= parent.end_token
|
|
and (child.start_token, child.end_token, child.phrase_id)
|
|
!= (parent.start_token, parent.end_token, parent.phrase_id)
|
|
)
|
|
|
|
|
|
def rank_match(match: HydratedPhraseMatch) -> tuple[float, float, int]:
|
|
"""Rank phrase matches by importance, confidence, then token count."""
|
|
return (match.importance, match.confidence, match.token_count)
|
|
|
|
|
|
def should_suppress(candidate: HydratedPhraseMatch, kept: HydratedPhraseMatch) -> bool:
|
|
"""Return whether an already-kept match should suppress a candidate."""
|
|
if not overlaps(candidate, kept):
|
|
return False
|
|
if candidate.canonical_id == kept.canonical_id:
|
|
return rank_match(kept) >= rank_match(candidate)
|
|
if is_inside(candidate, kept) and kept.suppress_children and not candidate.allow_nested:
|
|
return True
|
|
return not candidate.allow_nested and rank_match(kept) > rank_match(candidate)
|
|
|
|
|
|
def resolve_overlaps(matches: Sequence[HydratedPhraseMatch]) -> list[HydratedPhraseMatch]:
|
|
"""Resolve overlapping phrase matches without relying only on longest match."""
|
|
sorted_matches = sorted(
|
|
matches,
|
|
key=lambda match: (match.start_token, -match.token_count, -match.importance, -match.confidence),
|
|
)
|
|
kept: list[HydratedPhraseMatch] = []
|
|
for candidate in sorted_matches:
|
|
if any(should_suppress(candidate, existing) for existing in kept):
|
|
continue
|
|
kept.append(candidate)
|
|
return kept
|
|
|
|
|
|
def detect_protected_phrases_for_query(
|
|
session: Session,
|
|
query_text: str,
|
|
config: EbookSearchConfig,
|
|
*,
|
|
lookup: PhraseLookup | None = None,
|
|
book_id: int | None = None,
|
|
series_id: int | None = None,
|
|
) -> list[HydratedPhraseMatch]:
|
|
"""Run the full online protected-phrase query-detection pipeline."""
|
|
active_lookup = (
|
|
lookup if lookup is not None else load_phrase_lookup(session, config, book_id=book_id, series_id=series_id)
|
|
)
|
|
return resolve_overlaps(hydrate_matches(session, detect_phrase_candidates(query_text, active_lookup)))
|
|
|
|
|
|
def index_chunk_phrase_mentions_for_book(
|
|
session: Session,
|
|
book_id: int,
|
|
config: EbookSearchConfig,
|
|
*,
|
|
series_id: int | None = None,
|
|
lookup: PhraseLookup | None = None,
|
|
) -> int:
|
|
"""Rebuild chunk phrase mentions for all chunks in one book."""
|
|
active_lookup = (
|
|
lookup if lookup is not None else load_phrase_lookup(session, config, book_id=book_id, series_id=series_id)
|
|
)
|
|
session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.book_id == book_id))
|
|
chunks = session.scalars(select(EbookChunk).where(EbookChunk.source_id == book_id).order_by(EbookChunk.id))
|
|
count = 0
|
|
for chunk in chunks:
|
|
count += index_chunk_phrase_mentions(session, chunk, lookup=active_lookup)
|
|
session.flush()
|
|
logger.info("ebook_chunk_phrase_mentions_indexed book_id=%s mentions=%s", book_id, count)
|
|
return count
|
|
|
|
|
|
def index_chunk_phrase_mentions(session: Session, chunk: EbookChunk, *, lookup: PhraseLookup) -> int:
|
|
"""Store protected phrase mentions for one chunk."""
|
|
raw_matches = detect_phrase_candidates_in_text(chunk.text, lookup)
|
|
hydrated = resolve_overlaps(hydrate_matches(session, raw_matches))
|
|
for match in hydrated:
|
|
session.add(
|
|
EbookChunkPhraseMention(
|
|
chunk_id=chunk.id,
|
|
phrase_id=match.phrase_id,
|
|
book_id=match.book_id if match.book_id is not None else chunk.source_id,
|
|
series_id=match.series_id,
|
|
start_char=match.start_char if match.start_char is not None else 0,
|
|
end_char=match.end_char,
|
|
)
|
|
)
|
|
return len(hydrated)
|
|
|
|
|
|
def phrase_hit_counts_for_chunks(
|
|
session: Session,
|
|
*,
|
|
chunk_ids: Sequence[int],
|
|
phrase_ids: Sequence[int],
|
|
) -> dict[int, int]:
|
|
"""Return phrase-hit counts by chunk id using indexed chunk mentions."""
|
|
if not chunk_ids or not phrase_ids:
|
|
return {}
|
|
|
|
statement = (
|
|
select(
|
|
EbookChunkPhraseMention.chunk_id,
|
|
func.count(EbookChunkPhraseMention.phrase_id).label("phrase_hit_count"),
|
|
)
|
|
.where(
|
|
EbookChunkPhraseMention.chunk_id.in_(chunk_ids),
|
|
EbookChunkPhraseMention.phrase_id.in_(phrase_ids),
|
|
)
|
|
.group_by(EbookChunkPhraseMention.chunk_id)
|
|
)
|
|
return {int(row.chunk_id): int(row.phrase_hit_count) for row in session.execute(statement)}
|