"""Candidate phrase extraction and scoring for protected phrases.""" from __future__ import annotations import logging import re from collections import defaultdict from time import perf_counter from typing import TYPE_CHECKING, Protocol from yake import KeywordExtractor from python.ebook_search.protected_phrases.config import ( get_bad_ends, get_bad_starts, get_ignored_phrases, get_most_common_words, ) from python.ebook_search.protected_phrases.models import PhraseCandidate from python.ebook_search.protected_phrases.text_normalization import tokenize, tokenize_with_offsets if TYPE_CHECKING: from collections.abc import Iterable, Mapping, Sequence from python.ebook_search.config import EbookSearchConfig logger = logging.getLogger(__name__) 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", ) 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. Args: lan (str): Language code passed to YAKE. n (int): Maximum n-gram size to extract. dedupLim (float): Deduplication similarity threshold. top (int): Maximum number of keyphrases to return. Returns: YakeExtractor: The constructed keyword extractor. """ def strip_leading_articles(phrase_norm: str) -> str: """Remove one leading English article from a normalized phrase. Args: phrase_norm (str): Normalized phrase text to strip. Returns: str: The phrase with a single leading ``the``, ``a``, or ``an`` removed. """ 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. Args: phrase_text (str): Raw phrase text to normalize. config (EbookSearchConfig): Runtime phrase-tuning settings. min_tokens (int | None): Minimum token count override; defaults to ``config.phrase_min_tokens``. max_tokens (int | None): Maximum token count override; defaults to ``config.phrase_max_tokens``. strip_leading_article (bool): Whether to drop a single leading English article. Returns: tuple[str, str, int] | None: Display text, normalized phrase, and token count, or ``None`` when the phrase falls outside the token bounds or is ignored. """ 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. Args: text (str): Book text to slide n-gram windows over. config (EbookSearchConfig): Runtime phrase-tuning settings. Returns: dict[str, PhraseCandidate]: Candidates keyed by normalized phrase, with raw counts. """ 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. Args: book_text (str): Full book text to extract keyphrases from. config (EbookSearchConfig): Runtime phrase-tuning settings. top_k (int): Maximum number of YAKE keyphrases to request. Returns: dict[str, PhraseCandidate]: Candidates keyed by normalized phrase, with YAKE scores. """ extractor = KeywordExtractor(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 extract_spacy_candidates( book_text: str, nlp: SpacyLanguage, config: EbookSearchConfig, ) -> dict[str, PhraseCandidate]: """Extract spaCy named entities and noun chunks from one text block. Args: book_text (str): Text block to parse with spaCy. nlp (SpacyLanguage): Callable spaCy language pipeline. config (EbookSearchConfig): Runtime phrase-tuning settings. Returns: dict[str, PhraseCandidate]: Candidates keyed by normalized phrase from entities and noun chunks. """ 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. Args: original_text (str): Original-case book text to scan for capitalized runs. config (EbookSearchConfig): Runtime phrase-tuning settings. Returns: dict[str, PhraseCandidate]: Candidates keyed by normalized phrase from capitalized runs. """ 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. Args: metadata (Mapping[str, object] | None): Book metadata values, or ``None`` when unavailable. config (EbookSearchConfig): Runtime phrase-tuning settings. Returns: dict[str, PhraseCandidate]: Candidates keyed by normalized phrase from metadata values. """ 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. Args: *sources (Mapping[str, PhraseCandidate]): Candidate maps to combine, keyed by normalized phrase. Returns: dict[str, PhraseCandidate]: One merged candidate per 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. Args: existing (PhraseCandidate): Candidate mutated in place to absorb ``item``. item (PhraseCandidate): Candidate whose sources, counts, and scores are merged in. """ 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. Args: candidates (Mapping[str, PhraseCandidate]): Candidates to enrich, keyed by normalized phrase. chapters (Sequence[str]): Chapter-like text blocks used to count occurrences and spread. Returns: dict[str, PhraseCandidate]: Candidates with updated ``raw_count`` and ``chapter_count`` values. """ 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. Args: candidates (Mapping[str, PhraseCandidate]): Candidates to filter, keyed by normalized phrase. config (EbookSearchConfig): Runtime phrase-tuning settings. Returns: tuple[dict[str, PhraseCandidate], int, int, int]: The storable candidates followed by the counts dropped for being too short, too rare, and too common. """ 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. Args: config (EbookSearchConfig): Runtime phrase-tuning settings. Returns: int: The minimum raw occurrence count, never less than 1. """ 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. Args: phrase_norm (str): Normalized phrase text to inspect. Returns: bool: True when the phrase is non-empty and every token 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. Args: candidate (PhraseCandidate): Candidate to score. config (EbookSearchConfig): Runtime phrase-tuning settings. Returns: float: Combined score from sources, frequency, and length, less any penalties. """ 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. Args: phrase_norm (str): Normalized phrase text to inspect. Returns: bool: True when the first token is a known 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. Args: phrase_norm (str): Normalized phrase text to inspect. Returns: bool: True when the last token is a known 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. Args: candidate (PhraseCandidate): Candidate whose enabled sources are weighted. Returns: float: Summed weight of the candidate's enabled 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. Args: candidate (PhraseCandidate): Candidate whose counts are scored. config (EbookSearchConfig): Runtime phrase-tuning settings holding score thresholds. Returns: float: Summed weight for each frequency and chapter-spread threshold the candidate meets. """ 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. Args: candidate (PhraseCandidate): Candidate whose token count is scored. config (EbookSearchConfig): Runtime phrase-tuning settings holding the max token bound. Returns: float: Length-based score contribution, which may be negative for over- or under-length phrases. """ 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(normalized_book_text: str, phrase_norm: str, max_contexts: int = 5) -> list[str]: """Return normalized context snippets containing a candidate phrase. ``normalized_book_text`` is expected to already be ``normalize_text``-ed by the caller so the whole book is not re-normalized for every phrase. Args: normalized_book_text (str): Whole book text, already normalized, to search. phrase_norm (str): Normalized phrase to find contexts around. max_contexts (int): Maximum number of context snippets to return. Returns: list[str]: Up to ``max_contexts`` normalized snippets surrounding the phrase. """ contexts: list[str] = [] start = 0 while len(contexts) < max_contexts: index = normalized_book_text.find(phrase_norm, start) if index == -1: break left = max(0, index - 300) right = min(len(normalized_book_text), index + len(phrase_norm) + 300) contexts.append(normalized_book_text[left:right]) start = index + len(phrase_norm) return contexts def candidate_source_names(candidate: PhraseCandidate) -> list[str]: """Return enabled source names for an extracted candidate. Args: candidate (PhraseCandidate): Candidate whose enabled sources are listed. Returns: list[str]: Names of the extraction sources that produced the 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 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. Args: book_text (str): Full book text used for most extraction sources. chapters (Sequence[str]): Chapter-like text blocks used for spaCy and frequency counts. config (EbookSearchConfig): Runtime phrase-tuning settings. nlp (SpacyLanguage | None): Optional spaCy pipeline for entity and noun-chunk sources. metadata (Mapping[str, object] | None): Optional book metadata used as a candidate source. Returns: list[PhraseCandidate]: Scored candidates sorted best-first and capped per 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