436 lines
17 KiB
Python
436 lines
17 KiB
Python
"""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"(?<!\w)'[^'\n]+'(?!\w)")
|
|
REPORTING_VERB = (
|
|
r"(?:said|asked|replied|whispered|muttered|shouted|demanded|"
|
|
r"admitted|grumbled|called|cried|answered|warned|added)"
|
|
)
|
|
SPEECH_TAG = re.compile(
|
|
r"\b(he|she|they|it)\s+(?:(?:softly|quietly)\s+)?"
|
|
+ REPORTING_VERB + r"\b",
|
|
re.IGNORECASE,
|
|
)
|
|
PRONOUNS = frozenset(("he", "she", "they", "it"))
|
|
MAVERICK_TOKEN_LIMIT = 500
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AttributionStats:
|
|
total_paragraphs: int
|
|
labeled_paragraphs: int
|
|
unlabeled_paragraphs: int
|
|
expected_dialogue: int
|
|
expected_non_dialogue: int
|
|
correct_assignments: int
|
|
incorrect_assignments: int
|
|
missed_dialogue: int
|
|
extra_dialogue: int
|
|
correct_non_dialogue: int
|
|
predicted_on_unlabeled: int
|
|
|
|
|
|
def identify_dialogue(paragraph, word_start, italic_word_indexes, paragraph_index=None):
|
|
"""Detect quoted text or a fully italic or uppercase punctuated clause."""
|
|
|
|
if '"' in paragraph:
|
|
return True
|
|
|
|
return bool(_styled_dialogue_word_indexes(
|
|
paragraph, word_start, italic_word_indexes
|
|
))
|
|
|
|
|
|
def _styled_dialogue_word_indexes(paragraph, word_start, italic_word_indexes):
|
|
"""Find words in wholly italic or uppercase punctuated clauses."""
|
|
|
|
words = list(WORD.finditer(paragraph))
|
|
clauses = {}
|
|
for word_number, match in enumerate(words):
|
|
letters = "".join(character for character in match.group()
|
|
if character.isalpha())
|
|
clause_number = len(CLAUSE_PUNCTUATION.findall(paragraph[:match.start()]))
|
|
clause = clauses.setdefault(
|
|
clause_number,
|
|
{"italic": True, "uppercase": True, "last_start": match.start(),
|
|
"first_start": match.start(), "word_indexes": []},
|
|
)
|
|
clause["italic"] &= word_start + word_number in italic_word_indexes
|
|
clause["uppercase"] &= bool(letters) and letters.isupper()
|
|
clause["last_start"] = match.start()
|
|
clause["word_indexes"].append(word_start + word_number)
|
|
|
|
ordered = list(clauses.values())
|
|
dialogue_indexes = set()
|
|
for index, clause in enumerate(ordered):
|
|
boundary_end = (ordered[index + 1]["first_start"]
|
|
if index + 1 < len(ordered) else len(paragraph))
|
|
punctuated = CLAUSE_PUNCTUATION.search(
|
|
paragraph, clause["last_start"], boundary_end
|
|
) is not None
|
|
if punctuated and (clause["italic"] or clause["uppercase"]):
|
|
dialogue_indexes.update(clause["word_indexes"])
|
|
return dialogue_indexes
|
|
|
|
|
|
def dialogue_word_indexes(text, style_records):
|
|
"""Find quoted and styled dialogue words without resolving speakers."""
|
|
|
|
italic_indexes = {
|
|
index for index, flags, *_ in style_records if flags & 0x01
|
|
}
|
|
dialogue_indexes = set()
|
|
word_start = 0
|
|
in_double_quote = False
|
|
for raw_paragraph in PARAGRAPH_BREAK.split(text):
|
|
paragraph = raw_paragraph.strip()
|
|
if not paragraph:
|
|
continue
|
|
words = list(WORD.finditer(paragraph))
|
|
dialogue_indexes.update(_styled_dialogue_word_indexes(
|
|
paragraph, word_start, italic_indexes
|
|
))
|
|
single_quotes = list(PAIRED_SINGLE_QUOTES.finditer(paragraph))
|
|
continuation_quote = in_double_quote and paragraph.startswith('"')
|
|
for word_number, match in enumerate(words):
|
|
quoted_word = False
|
|
for position in range(match.start(), match.end()):
|
|
character = paragraph[position]
|
|
if character == '"':
|
|
quoted_word = True
|
|
if continuation_quote and position == 0:
|
|
continuation_quote = False
|
|
else:
|
|
in_double_quote = not in_double_quote
|
|
elif in_double_quote:
|
|
quoted_word = True
|
|
if quoted_word or any(
|
|
match.start() < quote.end() and match.end() > 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"(?<!\w)" + re.escape(name) + r"(?!\w)"
|
|
r"\s+(?:(?:softly|quietly)\s+)?" + REPORTING_VERB + r"\b",
|
|
re.IGNORECASE,
|
|
)
|
|
if pattern.search(text):
|
|
matches.add(name)
|
|
return next(iter(matches)) if len(matches) == 1 else None
|
|
|
|
|
|
@lru_cache(maxsize=2)
|
|
def load_maverick(device="cpu"):
|
|
"""Load the literary model with only its checkpoint metadata allowlisted."""
|
|
|
|
try:
|
|
from maverick import Maverick
|
|
from omegaconf.base import ContainerMetadata, Metadata
|
|
from omegaconf.dictconfig import DictConfig
|
|
from omegaconf.nodes import AnyNode
|
|
from torch.serialization import safe_globals
|
|
except ImportError as exc:
|
|
raise RuntimeError(
|
|
"Maverick attribution requires maverick-coref; install "
|
|
"book-processing/requirements.txt in the project venv"
|
|
) from exc
|
|
# PyTorch 2.6 loads checkpoints in weights-only mode. These are the
|
|
# non-tensor types in the published LitBank checkpoint's Hydra metadata.
|
|
with safe_globals([Any, dict, AnyNode, defaultdict, Metadata,
|
|
ContainerMetadata, DictConfig]):
|
|
return Maverick(hf_name_or_path="sapienzanlp/maverick-mes-litbank",
|
|
device=device)
|
|
|
|
|
|
def _align_tokens(text, tokens):
|
|
"""Align Maverick tokens against the input, preserving actual whitespace.
|
|
|
|
Maverick's plain-text char offsets assume one character between sentences,
|
|
which is not true for our blank-line-delimited chapters.
|
|
"""
|
|
|
|
spans = []
|
|
cursor = 0
|
|
for token in tokens:
|
|
start = text.find(token, cursor)
|
|
if start < 0:
|
|
raise ValueError(f"Cannot align Maverick token {token!r} after {cursor}")
|
|
end = start + len(token)
|
|
spans.append((start, end))
|
|
cursor = end
|
|
return spans
|
|
|
|
|
|
def _cluster_pronouns(text, result, speaker_names):
|
|
"""Map pronoun character starts to a name only for unambiguous clusters."""
|
|
|
|
token_spans = _align_tokens(text, result["tokens"])
|
|
names = tuple(name for name in speaker_names if name != "Other")
|
|
mappings = {}
|
|
for cluster in result["clusters_token_offsets"]:
|
|
mentions = [text[token_spans[start][0]:token_spans[end][1]]
|
|
for start, end in cluster]
|
|
matches = {
|
|
name for name in names for mention in mentions
|
|
if re.search(r"(?<!\w)" + re.escape(name) + r"(?!\w)",
|
|
mention, re.IGNORECASE)
|
|
}
|
|
if len(matches) != 1:
|
|
continue
|
|
name = next(iter(matches))
|
|
for start, end in cluster:
|
|
mention = text[token_spans[start][0]:token_spans[end][1]]
|
|
if mention.casefold() in PRONOUNS:
|
|
mappings[token_spans[start][0]] = name
|
|
return mappings
|
|
|
|
|
|
def _maverick_token_count(model, text):
|
|
"""Count the input tokens Maverick will send to its encoder."""
|
|
|
|
if hasattr(model, "token_count"):
|
|
return model.token_count(text)
|
|
tokens, _, _, _ = model.preprocess(text)
|
|
input_tokens = ["[SPEAKER_START]", "-", "[SPEAKER_END]", *tokens]
|
|
return len(model.tokenizer(input_tokens, add_special_tokens=True,
|
|
is_split_into_words=True)["input_ids"])
|
|
|
|
|
|
def _window_start(text, sentence_start, sentence_end, model):
|
|
"""Keep the target sentence and the most preceding words within budget."""
|
|
|
|
if _maverick_token_count(model, text[sentence_start:sentence_end]) > 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"(?<!\w)" + re.escape(entity) + r"(?!\w)",
|
|
paragraph, re.IGNORECASE)
|
|
scores[entity] = score + 5 if mentioned else max(0, score - 1)
|
|
|
|
|
|
def highest_relevance_except(scores, previous_speaker):
|
|
"""Choose a unique top scorer other than the previous dialogue entity."""
|
|
|
|
candidates = {entity: score for entity, score in scores.items()
|
|
if entity != previous_speaker}
|
|
if not candidates:
|
|
return "Other"
|
|
highest = max(candidates.values())
|
|
winners = [entity for entity, score in candidates.items()
|
|
if score == highest]
|
|
return winners[0] if len(winners) == 1 else "Other"
|
|
|
|
|
|
def evaluate_attributions(paragraphs, expected_by_index, italic_word_indexes,
|
|
dialogue_rule=identify_dialogue, speaker_names=(),
|
|
maverick_model=None, device="cpu"):
|
|
"""Apply pronoun, continuation, and relevance rules, then score labels."""
|
|
|
|
dialogue = [
|
|
dialogue_rule(
|
|
paragraph["text"], paragraph["word_start"], italic_word_indexes,
|
|
paragraph["index"],
|
|
)
|
|
for paragraph in paragraphs
|
|
]
|
|
named_speakers = {
|
|
paragraph["index"]: speaker
|
|
for paragraph, is_dialogue in zip(paragraphs, dialogue)
|
|
if is_dialogue
|
|
if (speaker := named_speech_tag_speaker(
|
|
paragraph["text"], speaker_names
|
|
)) is not None
|
|
}
|
|
targets = {
|
|
paragraph["index"]: speech_tag_pronouns(paragraph["text"])
|
|
for paragraph, is_dialogue in zip(paragraphs, dialogue)
|
|
if is_dialogue and paragraph["index"] not in named_speakers
|
|
}
|
|
initial, later = maverick_pronoun_speakers(
|
|
paragraphs, speaker_names, model=maverick_model, device=device,
|
|
target_pronouns=targets, return_stages=True,
|
|
)
|
|
def predict_with(resolved):
|
|
relevance = {name: 0 for name in dict.fromkeys(speaker_names)
|
|
if name != "Other"}
|
|
predictions = []
|
|
previous_speaker = None
|
|
for position, (paragraph, is_dialogue) in enumerate(zip(paragraphs, dialogue)):
|
|
update_entity_relevance(relevance, paragraph["text"])
|
|
if not is_dialogue:
|
|
prediction = None
|
|
elif paragraph["index"] in named_speakers:
|
|
prediction = named_speakers[paragraph["index"]]
|
|
elif (position > 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,
|
|
)
|