""" structural_segmenter.py ======================= Permissive segmenter for plays and novels. Handles smushed lines. Enhancements over v1: - Emotion tagging on dialogue blocks (lazily loaded EmotionAnalyzer) - Better character cue parsing (handles "ROMEO.", "LADY MACBETH", "FIRST WITCH") - French play/novel headings (ACTE, SCÈNE, CHAPITRE, PARTIE) - Stage direction parsing improved (handles both () and []) - Novel paragraph grouping: consecutive lines merged into paragraphs - language parameter propagated through for multilingual emotion analysis """ from __future__ import annotations import re import uuid from collections import namedtuple from typing import Any, Dict, List, Optional, Tuple import os from services.hf_inference_service import HFInferenceService # Initialize HF Inference Service _hf_inference = HFInferenceService(os.getenv("HF_API_TOKEN")) def _ai_is_heading(text: str) -> bool: """ Ask the cloud AI whether ``text`` is a chapter/section heading. Only called for short (≤ 80 char), isolated text blocks. """ if os.getenv("USE_HF_INFERENCE") == "1" and _hf_inference.api_token: try: return _hf_inference.is_heading(text) except Exception as e: print(f"⚠️ AI heading check error: {e}") pass return False # ── Data types ──────────────────────────────────────────────────────────────── Span = namedtuple("Span", ["text", "size", "flags", "page"]) # ── Helpers ─────────────────────────────────────────────────────────────────── def _uid() -> str: return str(uuid.uuid4())[:8] # ── Heading regexes ─────────────────────────────────────────────────────────── # English play headings _ACT_RE = re.compile(r"\bACT\s+([IVX]+|\d+)\b", re.IGNORECASE) _SCENE_RE = re.compile(r"\bSCENE\s+([IVX]+|\d+)\b", re.IGNORECASE) # Abbreviated scene (e.g. Folger Shakespeare "SC. 2") _SCENE_SC_RE = re.compile(r"\bSC\.\s*(\d+|[IVX]+)\b", re.IGNORECASE) # French play headings _ACTE_RE = re.compile(r"\bACTE\s+([IVX]+|\d+)\b", re.IGNORECASE) _SCENE_FR_RE = re.compile(r"\bSC[ÈE]NE\s+([IVX]+|\d+)\b", re.IGNORECASE) # Running-header detector: lines starting or ending with a page number, # or containing book title artifacts. _RUNNING_HEADER_RE = re.compile( r"^\d{1,4}\s+\S|" # "9 MACBETH..." r"^.*\s+\d{1,4}$|" # "...Before Breakfast 3" r"^[A-Z\s]{5,}\s+\d{1,4}$" # "CHARLOTTE'S WEB 4" ) # English novel headings _CHAPTER_RE = re.compile( r"\b(CHAPTER|PROLOGUE|EPILOGUE|PART|BOOK)\s+([\dIVX]+|[A-Z][a-z]+)?\b", re.IGNORECASE, ) # French novel headings _CHAPITRE_RE = re.compile(r"\bCHAPITRE\s+([\dIVX]+|[A-Z][a-z]+)?\b", re.IGNORECASE) _PARTIE_RE = re.compile(r"\bPARTIE\s+([\dIVX]+|[A-Z][a-z]+)?\b", re.IGNORECASE) _LIVRE_RE = re.compile(r"\bLIVRE\s+([\dIVX]+|[A-Z][a-z]+)?\b", re.IGNORECASE) # Character cue: ALL CAPS line (1–5 words), optional period/colon, <30 chars total # Handles: "ROMEO", "LADY MACBETH", "FIRST WITCH.", "THE GHOST:" _CUE_RE = re.compile( r"^[A-Z][A-Z\s'\.\d]{0,28}[A-Z]\.?\:?\s*$" ) # Stage directions _STAGE_INLINE_RE = re.compile(r"^\[.*\]$|^\(.*\)$") # Default neutral animation block _NEUTRAL_ANIM = { "expression": "neutral", "eyebrows": "neutral", "mouth": "closed", "eyes": "normal", "color_tint": "#f1f5f9", } # ── Segmenter ───────────────────────────────────────────────────────────────── class StructuralSegmenter: """ Converts a flat list of PyMuPDF block dicts into a hierarchical structure: - Plays: [Act → [Scene → [blocks]]] - Novels: [Chapter → [Section → paragraphs]] After segmentation, dialogue blocks in plays can be optionally enriched with emotion data via the EmotionAnalyzer singleton. """ # ── Public API ───────────────────────────────────────────────────────────── def segment( self, all_blocks: List[Dict[str, Any]], doc_type: str, language: str = "en", add_emotions: bool = True, ai_headings: Optional[List[Dict[str, str]]] = None, ) -> List[Dict[str, Any]]: """ Segment blocks into hierarchical units. Parameters ---------- all_blocks: List of PyMuPDF block dicts. doc_type: "play" | "novel" | "generic" language: "en" | "fr" add_emotions: If True, run EmotionAnalyzer on dialogue blocks (plays only). ai_headings: Optional list of metadata-inferred headings [{"title": "...", "text": "..."}]. """ self.ai_headings = ai_headings or [] tokens = self._tokenise(all_blocks, doc_type) if doc_type == "play": units = self._build_play_hierarchy(tokens) if add_emotions: self._enrich_play_emotions(units, language) return units if doc_type == "poem": return self._build_poem_hierarchy(tokens) return self._build_novel_hierarchy(tokens) # ── Heading threshold ─────────────────────────────────────────────────── @staticmethod def _compute_heading_threshold(spans: List[Span]) -> float: """Compute a font-size threshold above which text is likely a heading.""" if not spans: return 14.0 sizes = sorted(s.size for s in spans) n = len(sizes) median = sizes[n // 2] if n % 2 else (sizes[n // 2 - 1] + sizes[n // 2]) / 2 max_size = sizes[-1] return median + (max_size - median) * 0.5 # ── Tokenisation ────────────────────────────────────────────────────────── def _tokenise( self, blocks: List[Dict[str, Any]], doc_type: str ) -> List[Dict[str, Any]]: tokens: List[Dict[str, Any]] = [] for b in blocks: if b.get("type") != 0: continue for line_obj in b.get("lines", []): txt = "".join( s.get("text", "") for s in line_obj.get("spans", []) ).strip() if not txt: continue level, matched = self._match_heading(txt, doc_type) if matched: tokens.append({ "type": "heading", "level": level, "title": txt, "inferred": False, }) else: tokens.append({"type": "content", "text": txt}) return tokens def _match_heading(self, text: str, doc_type: str) -> Tuple[str, bool]: """ Check for heading patterns. Priority: AI-discovered headings -> Regex patterns. """ clean_text = self._normalize_heading(text) # 1. AI Anchor Match (highest priority) if hasattr(self, "ai_headings") and self.ai_headings: for ah in self.ai_headings: target = self._normalize_heading(ah.get("text", ah.get("title", ""))) if target and target in clean_text: # Treat as chapter-level for novels/poems, act/scene follows regex logic below return "chapter", True t = text[:30] # 2. Strong Regex Patterns (must check BEFORE running header as "CHAPTER 2" matches running header) if doc_type == "play": if _ACT_RE.search(t): return "act", True if _ACTE_RE.search(t): return "act", True if _SCENE_RE.search(t): return "scene", True if _SCENE_SC_RE.search(t): return "scene", True if _SCENE_FR_RE.search(t): return "scene", True else: if _CHAPTER_RE.search(t): return "chapter", True if _CHAPITRE_RE.search(t): return "chapter", True if _PARTIE_RE.search(t): return "chapter", True if _LIVRE_RE.search(t): return "chapter", True # 3. Skip running headers (page numbers, book titles at edges) if _RUNNING_HEADER_RE.match(text.strip()): return "none", False # 4. Mistral-7B fallback: ask the model for short, isolated text # that looks like it could be an unnamed heading (e.g. "The Storm", # "I.", roman numerals, numbered titles without the word CHAPTER). if len(text.strip()) <= 80 and _ai_is_heading(text.strip()): return "chapter", True return "none", False @staticmethod def _normalize_heading(text: str) -> str: """Lowercases and strips common punctuation/whitespace for comparison.""" return re.sub(r'[^a-zA-Z0-9]', '', text.lower()) # ── Play hierarchy ───────────────────────────────────────────────────────── def _build_play_hierarchy( self, tokens: List[Dict[str, Any]] ) -> List[Dict[str, Any]]: acts: List[Dict[str, Any]] = [] cur_act: Optional[Dict] = None cur_scene: Optional[Dict] = None content_buffer: List[str] = [] def flush(): if cur_scene is not None and content_buffer: cur_scene["blocks"].extend( self._parse_play_content(content_buffer) ) content_buffer.clear() for tok in tokens: if tok["type"] == "heading": if tok["level"] == "act": flush() if cur_act and cur_scene: cur_act["children"].append(cur_scene) if cur_act: acts.append(cur_act) cur_act = {"id": _uid(), "title": tok["title"], "children": []} cur_scene = None elif tok["level"] == "scene": flush() if cur_act and cur_scene: cur_act["children"].append(cur_scene) if cur_act is None: cur_act = {"id": _uid(), "title": "ACT I", "children": [], "inferred": True} cur_scene = {"id": _uid(), "title": tok["title"], "blocks": []} else: # content token if cur_act is None: continue if cur_scene is None: cur_scene = { "id": _uid(), "title": "Scene 1", "blocks": [], "inferred": True, } content_buffer.append(tok["text"]) # Flush final scene/act flush() if cur_act and cur_scene: cur_act["children"].append(cur_scene) if cur_act: acts.append(cur_act) return acts if acts else self._play_fallback(tokens) def _parse_play_content(self, lines: List[str]) -> List[Dict[str, Any]]: """ Convert flat text lines into structured block dicts. Block types: "dialogue" – character speech "stage_direction" – [Enter Juliet] or (aside) "narrative" – other text """ blocks: List[Dict[str, Any]] = [] for line in lines: stripped = line.strip() if not stripped: continue # Strip leading page numbers smushed with text (e.g. "11She's" -> "She's") stripped = re.sub(r'^\d{1,3}([A-Z])', r'\1', stripped) # Skip running headers if _RUNNING_HEADER_RE.match(stripped) or _RUNNING_HEADER_RE.search(stripped): if len(stripped) < 50: # Headers are usually short continue # Stage direction: entire line wrapped in () or [] if _STAGE_INLINE_RE.match(stripped): blocks.append({ "type": "stage_direction", "character": None, "content": stripped.strip("[]()").strip(), }) continue # Character cue (ALL CAPS, ≤6 words, short line) word_count = len(stripped.split()) if _CUE_RE.match(stripped) and word_count <= 6: cue = stripped.rstrip(":. \t") blocks.append({ "type": "dialogue", "character": cue, "content": "", "emotion": "neutral", "intensity": 0.5, "anim": dict(_NEUTRAL_ANIM), }) continue # Append to last dialogue block if mid-speech if blocks and blocks[-1]["type"] == "dialogue": # Skip bare page numbers (e.g. "5", "12") slipping into speech if re.match(r"^\d{1,4}$", stripped): continue sep = " " if blocks[-1]["content"] else "" blocks[-1]["content"] += sep + stripped continue # Narrative line blocks.append({"type": "narrative", "content": stripped}) return blocks def _play_fallback(self, tokens: List[Dict[str, Any]]) -> List[Dict[str, Any]]: """Single-unit fallback when no act/scene headings found.""" text_lines = [t["text"] for t in tokens if t["type"] == "content"] return [{ "id": _uid(), "title": "The Play", "children": [{ "id": _uid(), "title": "Content", "blocks": self._parse_play_content(text_lines), }], }] # ── Novel hierarchy ──────────────────────────────────────────────────────── def _build_novel_hierarchy( self, tokens: List[Dict[str, Any]] ) -> List[Dict[str, Any]]: chapters: List[Dict[str, Any]] = [] MAX_CHUNK_WORDS = 250 def new_chapter(title: str) -> Dict[str, Any]: return { "id": _uid(), "title": title, "children": [], # We will add chunks here "content": "", } def new_chunk(title: str = "") -> Dict[str, Any]: return { "id": _uid(), "title": title, "paragraphs": [], "content": "", "word_count": 0 } cur_chapter = new_chapter("Beginning") cur_chunk = new_chunk("Introduction") for tok in tokens: if tok["type"] == "heading" and tok["level"] == "chapter": # Flush current chunk and chapter if cur_chunk["paragraphs"]: cur_chapter["children"].append(cur_chunk) if cur_chapter["children"]: self._finalize_chapter(cur_chapter) chapters.append(cur_chapter) cur_chapter = new_chapter(tok["title"]) cur_chunk = new_chunk() elif tok["type"] == "content": text = tok["text"] words = len(text.split()) # If adding this would explode the chunk, flush first if cur_chunk["word_count"] > 0 and (cur_chunk["word_count"] + words) > MAX_CHUNK_WORDS: cur_chapter["children"].append(cur_chunk) cur_chunk = new_chunk() cur_chunk["paragraphs"].append(text) cur_chunk["word_count"] += words # Final flush if cur_chunk["paragraphs"]: cur_chapter["children"].append(cur_chunk) if cur_chapter["children"]: self._finalize_chapter(cur_chapter) chapters.append(cur_chapter) return chapters @staticmethod def _finalize_chapter(chapter: Dict[str, Any]): """Finalize all chunks in the chapter and set overall content.""" all_content = [] for i, chunk in enumerate(chapter["children"]): chunk["content"] = "\n\n".join(chunk.get("paragraphs", [])) # If no specific title, give it a sequence title if not chunk["title"]: chunk["title"] = f"Part {i+1}" all_content.append(chunk["content"]) chapter["content"] = "\n\n".join(all_content) # ── Poem hierarchy ───────────────────────────────────────────────────────── def _build_poem_hierarchy( self, tokens: List[Dict[str, Any]] ) -> List[Dict[str, Any]]: """ Build stanza-based hierarchy for poems. Groups consecutive content lines into stanzas, separated by blank gaps or heading tokens. """ poems: List[Dict[str, Any]] = [] current_poem: Optional[Dict] = None current_stanza: List[str] = [] stanza_num = 0 def flush_stanza(): nonlocal stanza_num if current_poem is not None and current_stanza: stanza_num += 1 current_poem["children"].append({ "id": _uid(), "title": f"Stanza {stanza_num}", "blocks": [ {"type": "verse_line", "content": line} for line in current_stanza ], "content": "\n".join(current_stanza), }) current_stanza.clear() for tok in tokens: if tok["type"] == "heading": flush_stanza() if current_poem and current_poem["children"]: poems.append(current_poem) current_poem = {"id": _uid(), "title": tok["title"], "children": []} stanza_num = 0 else: text = tok.get("text", "").strip() if not text: # Empty line = stanza break flush_stanza() continue if current_poem is None: current_poem = {"id": _uid(), "title": "The Poem", "children": []} current_stanza.append(text) # Flush remaining flush_stanza() if current_poem and current_poem["children"]: poems.append(current_poem) if not poems: # Fallback: treat all content as one poem all_lines = [t["text"] for t in tokens if t["type"] == "content"] stanzas = [] current: List[str] = [] for line in all_lines: if not line.strip(): if current: stanzas.append(current) current = [] else: current.append(line) if current: stanzas.append(current) children = [ { "id": _uid(), "title": f"Stanza {i + 1}", "blocks": [{"type": "verse_line", "content": l} for l in st], "content": "\n".join(st), } for i, st in enumerate(stanzas) ] if stanzas else [{ "id": _uid(), "title": "Content", "blocks": [{"type": "verse_line", "content": l} for l in all_lines], "content": "\n".join(all_lines), }] poems = [{ "id": _uid(), "title": "The Poem", "children": children, }] return poems # ── Emotion enrichment ──────────────────────────────────────────────────── def _enrich_play_emotions( self, acts: List[Dict[str, Any]], language: str = "en" ): """ Walk Act → Scene → Block tree and add emotion data to all dialogue blocks. """ try: from .emotion_analyzer import get_emotion_analyzer ea = get_emotion_analyzer() except Exception as e: print(f"⚠️ StructuralSegmenter: emotion enrichment skipped ({e})") return for act in acts: for scene in act.get("children", []): blocks = scene.get("blocks", []) try: ea.enrich_dialogue_blocks(blocks, language=language) except Exception as exc: print( f"⚠️ Emotion error in scene " f"'{scene.get('title', '?')}': {exc}" )