Richie
7b97f31fec
- 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.
275 lines
9.6 KiB
Python
275 lines
9.6 KiB
Python
"""Book-level orchestration for candidate n-gram generation and recalculation."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from time import perf_counter
|
|
from typing import TYPE_CHECKING
|
|
|
|
from sqlalchemy import select
|
|
|
|
from python.ebook_search.protected_phrases.extraction import extract_phrase_candidates_for_book
|
|
from python.ebook_search.protected_phrases.models import (
|
|
BookCandidateResult,
|
|
PhraseCandidateGenerationResult,
|
|
PhraseRecalculationResult,
|
|
)
|
|
from python.ebook_search.protected_phrases.store import (
|
|
delete_phrase_data_for_book,
|
|
load_book_chapter_texts,
|
|
metadata_for_source,
|
|
prune_unstorable_unjudged_candidate_phrases,
|
|
save_candidate_to_db,
|
|
)
|
|
from python.orm.richie import EbookSource
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Mapping, Sequence
|
|
|
|
from sqlalchemy.orm import Session
|
|
|
|
from python.ebook_search.config import EbookSearchConfig
|
|
from python.ebook_search.protected_phrases.extraction import SpacyLanguage
|
|
from python.orm.richie import EbookCandidatePhrase
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def generate_candidate_phrases_for_books(
|
|
session: Session,
|
|
config: EbookSearchConfig,
|
|
*,
|
|
nlp: SpacyLanguage | None = None,
|
|
) -> PhraseCandidateGenerationResult:
|
|
"""Create or refresh candidate phrases for indexed books without calling the LLM judge.
|
|
|
|
Args:
|
|
session (Session): Active database session.
|
|
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
|
nlp (SpacyLanguage | None): Optional spaCy pipeline for entity and noun-chunk sources.
|
|
|
|
Returns:
|
|
PhraseCandidateGenerationResult: Per-corpus counts of books seen, built, and candidates stored.
|
|
"""
|
|
sources = session.scalars(select(EbookSource).order_by(EbookSource.id)).all()
|
|
books_seen = len(sources)
|
|
logger.info(
|
|
"ebook_candidate_phrase_generation_start books_seen=%s min_tokens=%s max_tokens=%s max_candidates_per_book=%s",
|
|
books_seen,
|
|
config.phrase_min_tokens,
|
|
config.phrase_max_tokens,
|
|
config.protected_phrase_max_candidates_per_book,
|
|
)
|
|
|
|
outcomes = [generate_candidates_for_source(session, source, config, nlp=nlp) for source in sources]
|
|
|
|
result = PhraseCandidateGenerationResult(
|
|
books_seen=books_seen,
|
|
books_built=sum(1 for outcome in outcomes if outcome.built),
|
|
candidate_phrases=sum(outcome.candidates for outcome in outcomes),
|
|
)
|
|
logger.info(
|
|
"ebook_candidate_phrase_generation_complete books_seen=%s books_built=%s candidate_total=%s",
|
|
result.books_seen,
|
|
result.books_built,
|
|
result.candidate_phrases,
|
|
)
|
|
return result
|
|
|
|
|
|
def generate_candidates_for_source(
|
|
session: Session,
|
|
source: EbookSource,
|
|
config: EbookSearchConfig,
|
|
*,
|
|
nlp: SpacyLanguage | None = None,
|
|
) -> BookCandidateResult:
|
|
"""Generate and store candidate phrases for one book, managing its own transaction.
|
|
|
|
Commits on success; rolls back and re-raises on error so callers stop the backfill.
|
|
|
|
Args:
|
|
session (Session): Active database session.
|
|
source (EbookSource): Indexed book to generate candidates for.
|
|
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
|
nlp (SpacyLanguage | None): Optional spaCy pipeline for entity and noun-chunk sources.
|
|
|
|
Returns:
|
|
BookCandidateResult: Candidate count and whether the book was committed.
|
|
"""
|
|
book_started_at = perf_counter()
|
|
logger.info(
|
|
"ebook_candidate_phrase_generation_book_start source_id=%s title=%r",
|
|
source.id,
|
|
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)
|
|
return BookCandidateResult()
|
|
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", source.id)
|
|
raise
|
|
|
|
logger.info(
|
|
"ebook_candidate_phrase_generation_book_committed source_id=%s candidates=%s duration_ms=%.1f",
|
|
source.id,
|
|
len(candidates),
|
|
(perf_counter() - book_started_at) * 1000,
|
|
)
|
|
return BookCandidateResult(candidates=len(candidates), built=True)
|
|
|
|
|
|
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.
|
|
|
|
Args:
|
|
session (Session): Active database session.
|
|
source (EbookSource): Indexed book to recalculate.
|
|
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
|
nlp (SpacyLanguage | None): Optional spaCy pipeline for entity and noun-chunk sources.
|
|
|
|
Returns:
|
|
PhraseRecalculationResult: Deleted-row counts and the number of candidates regenerated.
|
|
"""
|
|
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 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.
|
|
|
|
Args:
|
|
session (Session): Active database session.
|
|
book_id (int): Book the candidates belong to.
|
|
series_id (int | None): Series scope for the stored candidates.
|
|
book_text (str): Full book text used for extraction.
|
|
chapters (Sequence[str]): Chapter-like text blocks used for 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[EbookCandidatePhrase]: The stored candidate phrase rows.
|
|
"""
|
|
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
|