Richie
4861f58f27
- Introduced dataclasses for phrase candidates, judgments, and matches in `models.py`. - Implemented database operations for candidate and protected phrases in `store.py`, including loading, saving, and deleting phrases. - Enhanced text normalization functions in `text_normalization.py` with detailed docstrings. - Refactored search functionality to utilize new models and methods for detecting protected phrases.
488 lines
18 KiB
Python
488 lines
18 KiB
Python
"""Runtime protected-phrase matching and chunk mention indexing."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from collections import defaultdict
|
|
from typing import TYPE_CHECKING
|
|
|
|
from sqlalchemy import and_, delete, func, or_, select
|
|
|
|
from python.ebook_search.protected_phrases.config import get_ignored_phrases
|
|
from python.ebook_search.protected_phrases.models import (
|
|
ChunkPhraseHit,
|
|
HydratedPhraseMatch,
|
|
PhraseLookup,
|
|
PhraseMatch,
|
|
)
|
|
from python.ebook_search.protected_phrases.text_normalization import tokenize_with_offsets
|
|
from python.orm.richie import (
|
|
EbookChunk,
|
|
EbookChunkPhraseMention,
|
|
EbookPhraseAlias,
|
|
EbookProtectedPhrase,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Iterator, Sequence
|
|
|
|
from sqlalchemy.orm import Session
|
|
|
|
from python.ebook_search.config import EbookSearchConfig
|
|
from python.ebook_search.protected_phrases.text_normalization import NormalizedToken
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
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.
|
|
|
|
Args:
|
|
session (Session): Active database session.
|
|
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
|
book_id (int | None): Optional book scope to restrict loaded phrases.
|
|
series_id (int | None): Optional series scope to restrict loaded phrases.
|
|
|
|
Returns:
|
|
PhraseLookup: Normalized phrase and alias maps with the token-window bounds to test.
|
|
"""
|
|
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.
|
|
|
|
Args:
|
|
book_id (int | None): Optional book scope to include alongside global phrases.
|
|
series_id (int | None): Optional series scope to include alongside global phrases.
|
|
|
|
Returns:
|
|
object | None: A combined SQLAlchemy filter clause, or ``None`` when no scope is given.
|
|
"""
|
|
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.
|
|
|
|
Args:
|
|
tokens_ (Sequence[str]): Normalized query tokens.
|
|
min_n (int): Smallest window size to yield.
|
|
max_n (int): Largest window size to yield, capped at the token count.
|
|
|
|
Yields:
|
|
tuple[str, int, int]: Normalized window text with its start and end token indices.
|
|
"""
|
|
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.
|
|
|
|
Args:
|
|
query_text (str): User query text to scan.
|
|
lookup (PhraseLookup): In-memory phrase and alias lookup maps.
|
|
|
|
Returns:
|
|
list[PhraseMatch]: Unhydrated phrase matches found in the query.
|
|
"""
|
|
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.
|
|
|
|
Args:
|
|
text (str): Arbitrary text, such as a chunk, to scan.
|
|
lookup (PhraseLookup): In-memory phrase and alias lookup maps.
|
|
|
|
Returns:
|
|
list[PhraseMatch]: Unhydrated phrase matches found in the text.
|
|
"""
|
|
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.
|
|
|
|
Args:
|
|
tokens_ (Sequence[NormalizedToken]): Normalized tokens with character offsets.
|
|
lookup (PhraseLookup): In-memory phrase and alias lookup maps.
|
|
|
|
Returns:
|
|
list[PhraseMatch]: Deduplicated unhydrated phrase matches with token and character spans.
|
|
"""
|
|
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.
|
|
|
|
Args:
|
|
session (Session): Active database session.
|
|
matches (Sequence[PhraseMatch]): Unhydrated matches to enrich.
|
|
|
|
Returns:
|
|
list[HydratedPhraseMatch]: Matches with protected-phrase metadata attached.
|
|
"""
|
|
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.
|
|
|
|
Args:
|
|
first (HydratedPhraseMatch): First match to compare.
|
|
second (HydratedPhraseMatch): Second match to compare.
|
|
|
|
Returns:
|
|
bool: True when the two token spans share at least one token position.
|
|
"""
|
|
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.
|
|
|
|
Args:
|
|
child (HydratedPhraseMatch): Candidate nested match.
|
|
parent (HydratedPhraseMatch): Candidate enclosing match.
|
|
|
|
Returns:
|
|
bool: True when ``child`` lies within ``parent`` and is not the same span.
|
|
"""
|
|
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.
|
|
|
|
Args:
|
|
match (HydratedPhraseMatch): Match to build a sort key for.
|
|
|
|
Returns:
|
|
tuple[float, float, int]: A comparable key of importance, confidence, and 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.
|
|
|
|
Args:
|
|
candidate (HydratedPhraseMatch): Match being considered for keeping.
|
|
kept (HydratedPhraseMatch): Match already kept that may suppress the candidate.
|
|
|
|
Returns:
|
|
bool: True when the candidate should be dropped in favor of the kept match.
|
|
"""
|
|
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.
|
|
|
|
Args:
|
|
matches (Sequence[HydratedPhraseMatch]): Hydrated matches that may overlap.
|
|
|
|
Returns:
|
|
list[HydratedPhraseMatch]: The kept, non-suppressed matches.
|
|
"""
|
|
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.
|
|
|
|
Args:
|
|
session (Session): Active database session.
|
|
query_text (str): User query text to detect phrases in.
|
|
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
|
lookup (PhraseLookup | None): Optional preloaded lookup; loaded on demand when ``None``.
|
|
book_id (int | None): Optional book scope for lookup loading.
|
|
series_id (int | None): Optional series scope for lookup loading.
|
|
|
|
Returns:
|
|
list[HydratedPhraseMatch]: Hydrated, overlap-resolved phrase matches for the query.
|
|
"""
|
|
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.
|
|
|
|
Args:
|
|
session (Session): Active database session.
|
|
book_id (int): Book whose chunk mentions are rebuilt.
|
|
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
|
series_id (int | None): Optional series scope for lookup loading.
|
|
lookup (PhraseLookup | None): Optional preloaded lookup; loaded on demand when ``None``.
|
|
|
|
Returns:
|
|
int: Total number of chunk phrase mentions indexed for the 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.
|
|
|
|
Args:
|
|
session (Session): Active database session.
|
|
chunk (EbookChunk): Chunk whose text is scanned for phrase mentions.
|
|
lookup (PhraseLookup): In-memory phrase and alias lookup maps.
|
|
|
|
Returns:
|
|
int: Number of phrase mentions stored for the 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_hits_for_chunks(
|
|
session: Session,
|
|
*,
|
|
chunk_ids: Sequence[int],
|
|
phrase_ids: Sequence[int],
|
|
) -> dict[int, tuple[ChunkPhraseHit, ...]]:
|
|
"""Return matched protected phrases with mention counts by chunk id using indexed chunk mentions.
|
|
|
|
Args:
|
|
session (Session): Active database session.
|
|
chunk_ids (Sequence[int]): Chunk ids to look up mentions for.
|
|
phrase_ids (Sequence[int]): Protected phrase ids to restrict the results to.
|
|
|
|
Returns:
|
|
dict[int, tuple[ChunkPhraseHit, ...]]: Phrase hits per chunk id, ordered by mention count.
|
|
"""
|
|
if not chunk_ids or not phrase_ids:
|
|
return {}
|
|
|
|
mention_count = func.count(EbookChunkPhraseMention.phrase_id).label("mention_count")
|
|
statement = (
|
|
select(
|
|
EbookChunkPhraseMention.chunk_id,
|
|
EbookProtectedPhrase.id.label("phrase_id"),
|
|
EbookProtectedPhrase.phrase_text,
|
|
mention_count,
|
|
)
|
|
.join(EbookProtectedPhrase, EbookProtectedPhrase.id == EbookChunkPhraseMention.phrase_id)
|
|
.where(
|
|
EbookChunkPhraseMention.chunk_id.in_(chunk_ids),
|
|
EbookChunkPhraseMention.phrase_id.in_(phrase_ids),
|
|
)
|
|
.group_by(EbookChunkPhraseMention.chunk_id, EbookProtectedPhrase.id, EbookProtectedPhrase.phrase_text)
|
|
.order_by(EbookChunkPhraseMention.chunk_id, mention_count.desc(), EbookProtectedPhrase.phrase_text)
|
|
)
|
|
hits: defaultdict[int, list[ChunkPhraseHit]] = defaultdict(list)
|
|
for row in session.execute(statement):
|
|
hits[int(row.chunk_id)].append(
|
|
ChunkPhraseHit(
|
|
phrase_id=int(row.phrase_id),
|
|
phrase_text=str(row.phrase_text),
|
|
mention_count=int(row.mention_count),
|
|
)
|
|
)
|
|
return {chunk_id: tuple(chunk_hits) for chunk_id, chunk_hits in hits.items()}
|
|
|
|
|
|
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.
|
|
|
|
Args:
|
|
session (Session): Active database session.
|
|
chunk_ids (Sequence[int]): Chunk ids to count mentions for.
|
|
phrase_ids (Sequence[int]): Protected phrase ids to restrict the counts to.
|
|
|
|
Returns:
|
|
dict[int, int]: Total mention count per chunk id.
|
|
"""
|
|
hits = phrase_hits_for_chunks(session, chunk_ids=chunk_ids, phrase_ids=phrase_ids)
|
|
return {chunk_id: sum(hit.mention_count for hit in chunk_hits) for chunk_id, chunk_hits in hits.items()}
|