"""Dialogue detection and speaker attribution rules for chapter evaluation.""" import re from collections import defaultdict from dataclasses import dataclass from functools import lru_cache from typing import Any WORD = re.compile(r"\S+") CLAUSE_PUNCTUATION = re.compile(r"[.!?;,:—–]") PARAGRAPH_BREAK = re.compile(r"\r?\n[ \t]*\r?\n+") PAIRED_SINGLE_QUOTES = re.compile(r"(? quote.start() for quote in single_quotes ): dialogue_indexes.add(word_start + word_number) word_start += len(words) return dialogue_indexes def speech_tag_pronouns(paragraph): """Find reporting-clause pronouns outside double-quoted speech.""" return [match.start(1) for match in SPEECH_TAG.finditer(paragraph) if paragraph[:match.start(1)].count('"') % 2 == 0] def named_speech_tag_speaker(paragraph, speaker_names): """Return a unique named speaker followed by a reporting verb outside quotes.""" outside_quotes = [] quoted = False for character in paragraph: if character == '"': quoted = not quoted outside_quotes.append(" ") else: outside_quotes.append(" " if quoted else character) text = "".join(outside_quotes) matches = set() for name in speaker_names: if name == "Other": continue pattern = re.compile( r"(? MAVERICK_TOKEN_LIMIT: raise ValueError("A reporting-clause sentence exceeds Maverick's token budget") starts = sorted({0, sentence_start, *(match.start() for match in WORD.finditer(text[:sentence_start]))}) left, right = 0, len(starts) - 1 while left < right: middle = (left + right) // 2 if _maverick_token_count(model, text[starts[middle]:sentence_end]) <= MAVERICK_TOKEN_LIMIT: right = middle else: left = middle + 1 while _maverick_token_count(model, text[starts[left]:sentence_end]) > MAVERICK_TOKEN_LIMIT: left += 1 while left and _maverick_token_count(model, text[starts[left - 1]:sentence_end]) <= MAVERICK_TOKEN_LIMIT: left -= 1 return starts[left] def _maverick_resolutions(paragraphs, speaker_names, model, target_pronouns): """Resolve each tag in its sentence and revisit earlier tags in later runs.""" from nltk.data import load as load_nltk_resource sentence_tokenizer = load_nltk_resource("tokenizers/punkt/english.pickle") text = "\n\n".join(paragraph["text"] for paragraph in paragraphs) offset = 0 runs = [] tag_positions = {} for paragraph in paragraphs: index = paragraph["index"] tags = target_pronouns.get(index, ()) if tags: tag_positions[index] = tuple(offset + tag for tag in tags) spans = list(sentence_tokenizer.span_tokenize(paragraph["text"])) for start, end in spans: if any(start <= tag < end for tag in tags): runs.append((index, offset + start, offset + end)) offset += len(paragraph["text"]) + 2 initial = {} later = {} for current_index, sentence_start, sentence_end in runs: start = _window_start(text, sentence_start, sentence_end, model) mappings = _cluster_pronouns( text[start:sentence_end], model.predict(text[start:sentence_end]), speaker_names, ) for index, positions in tag_positions.items(): if index in initial or index in later: continue if not all(start <= position < sentence_end for position in positions): continue names = {mappings.get(position - start) for position in positions} if len(names) == 1 and None not in names: (initial if index == current_index else later)[index] = next(iter(names)) return initial, later def maverick_pronoun_speakers(paragraphs, speaker_names, model=None, device="cpu", target_pronouns=None, return_stages=False): """Resolve reporting-clause pronouns to a supplied cast of entity names. `model` accepts an object with Maverick's `predict(text)` API for tests. A missing or ambiguous mapping is intentionally omitted. """ if target_pronouns is None: target_pronouns = { paragraph["index"]: speech_tag_pronouns(paragraph["text"]) for paragraph in paragraphs } if not speaker_names or not any(target_pronouns.values()): return ({}, {}) if return_stages else {} if model is None: model = load_maverick(device) initial, later = _maverick_resolutions( paragraphs, speaker_names, model, target_pronouns, ) return (initial, later) if return_stages else {**later, **initial} def quotation_continues(previous_paragraph, paragraph): """Detect a new paragraph beginning inside the previous open quotation.""" return (previous_paragraph.count('"') % 2 == 1 and paragraph.startswith('"')) def update_entity_relevance(scores, paragraph): """Add five once for a named reference, otherwise decay one to zero.""" for entity, score in scores.items(): mentioned = re.search(r"(? 0 and dialogue[position - 1] and quotation_continues(paragraphs[position - 1]["text"], paragraph["text"])): prediction = predictions[-1] elif paragraph["index"] in resolved: prediction = resolved[paragraph["index"]] else: prediction = highest_relevance_except(relevance, previous_speaker) predictions.append(prediction) if is_dialogue: previous_speaker = prediction return predictions first_predictions = predict_with(initial) upgrades = {index: name for index, name in later.items() if first_predictions[index] == "Other"} predictions = predict_with({**initial, **upgrades}) if upgrades else first_predictions counts = {key: 0 for key in ( "correct_assignments", "incorrect_assignments", "missed_dialogue", "extra_dialogue", "correct_non_dialogue", )} predicted_on_unlabeled = 0 for index, prediction in enumerate(predictions): if index not in expected_by_index: predicted_on_unlabeled += prediction is not None continue expected = expected_by_index[index] if expected is None: key = "correct_non_dialogue" if prediction is None else "extra_dialogue" elif prediction is None: key = "missed_dialogue" elif prediction == expected: key = "correct_assignments" else: key = "incorrect_assignments" counts[key] += 1 expected_dialogue = sum(entity is not None for entity in expected_by_index.values()) return predictions, AttributionStats( total_paragraphs=len(paragraphs), labeled_paragraphs=len(expected_by_index), unlabeled_paragraphs=len(paragraphs) - len(expected_by_index), expected_dialogue=expected_dialogue, expected_non_dialogue=len(expected_by_index) - expected_dialogue, predicted_on_unlabeled=predicted_on_unlabeled, **counts, )