import json import os import re import unittest from pathlib import Path from attribution import (evaluate_attributions, maverick_pronoun_speakers, quotation_continues, speech_tag_pronouns, update_entity_relevance) from label_dialogue import paragraphs_with_word_indexes, read_italic_word_indexes FIXTURE_ROOT = Path(__file__).parent / "fixtures" / "attribution" FIXTURES = ( ("Prologue", "prologue", "Prologue- To Live"), ("Chapter 1", "chapter-1", "Chapter 1"), ) class StubMaverick: """Return Maverick-shaped clusters for named token occurrences.""" def __init__(self, clusters): self.clusters = clusters self.calls = 0 self.inputs = [] def token_count(self, text): return len(re.findall(r"\w+|[^\w\s]", text)) + 5 def predict(self, text): self.calls += 1 self.inputs.append(text) tokens = re.findall(r"\w+|[^\w\s]", text) offsets = [] for cluster in self.clusters: offsets.append([ (next(index for index, token in enumerate(tokens) if token == word and sum(t == word for t in tokens[:index]) == occurrence),) * 2 for word, occurrence in cluster ]) return {"tokens": tokens, "clusters_token_offsets": offsets} class LateNameMaverick(StubMaverick): """Link an earlier pronoun once a later window contains its name.""" def __init__(self): super().__init__([]) def predict(self, text): self.calls += 1 self.inputs.append(text) tokens = re.findall(r"\w+|[^\w\s]", text) clusters = [] if "Syl" in tokens and "she" in tokens: clusters.append([(tokens.index("Syl"),) * 2, (tokens.index("she"),) * 2]) return {"tokens": tokens, "clusters_token_offsets": clusters} class AttributionEvaluationTests(unittest.TestCase): def test_maverick_resolves_speech_tag_pronoun(self): text = 'Syl entered.\n\n"Hello," she said.' paragraphs = paragraphs_with_word_indexes(text) model = StubMaverick([[("Syl", 0), ("she", 0)]]) predictions, stats = evaluate_attributions( paragraphs, {1: "Syl"}, set(), speaker_names=("Syl", "Kaladin"), maverick_model=model, ) self.assertEqual(predictions, [None, "Syl"]) self.assertEqual(stats.correct_assignments, 1) self.assertEqual(model.calls, 1) def test_maverick_window_keeps_sentence_and_at_most_500_tokens(self): text = (" ".join(f"word{i}" for i in range(600)) + '\n\n"Hello," she said.') paragraphs = paragraphs_with_word_indexes(text) model = StubMaverick([]) maverick_pronoun_speakers(paragraphs, ("Syl",), model=model) self.assertEqual(model.calls, 1) self.assertTrue(model.inputs[0].endswith('"Hello," she said.')) self.assertLessEqual(model.token_count(model.inputs[0]), 500) self.assertGreater(model.token_count(model.inputs[0]), 480) def test_later_window_upgrades_only_other_to_named_entity(self): text = ('"Hello," she said.\n\nSyl arrived.\n\n' '"Goodbye," he said.') paragraphs = paragraphs_with_word_indexes(text) model = LateNameMaverick() predictions, _ = evaluate_attributions( paragraphs, {}, set(), speaker_names=("Syl", "Kaladin"), maverick_model=model, ) self.assertEqual(model.calls, 2) self.assertEqual(predictions[0], "Syl") def test_quoted_pronoun_and_unresolved_tag_do_not_assign_speaker(self): text = 'Syl entered.\n\n"She said goodbye," he said.' paragraphs = paragraphs_with_word_indexes(text) model = StubMaverick([[("Syl", 0), ("She", 0)]]) mappings = maverick_pronoun_speakers( paragraphs, ("Syl",), model=model, ) self.assertEqual(speech_tag_pronouns(paragraphs[1]["text"]), [20]) self.assertEqual(mappings, {}) def test_ambiguous_cluster_and_partly_unresolved_tags_stay_other(self): text = 'Syl met Kaladin.\n\n"Hello," she said, and he replied.' paragraphs = paragraphs_with_word_indexes(text) ambiguous = StubMaverick([[("Syl", 0), ("Kaladin", 0), ("she", 0)]]) partial = StubMaverick([[("Syl", 0), ("she", 0)]]) for model in (ambiguous, partial): with self.subTest(model=model): mappings = maverick_pronoun_speakers( paragraphs, ("Syl", "Kaladin"), model=model, ) self.assertEqual(mappings, {}) def test_open_quote_continues_across_paragraphs(self): text = ('Syl stood nearby.\n\n' '"Listen," she said. "This is not the end.\n\n' '"Still talking.\n\n' '"Now done."\n\n' '"Next?"') paragraphs = paragraphs_with_word_indexes(text) model = StubMaverick([[("Syl", 0), ("she", 0)]]) predictions, _ = evaluate_attributions( paragraphs, {}, set(), speaker_names=("Syl", "Kaladin"), maverick_model=model, ) self.assertTrue(quotation_continues(paragraphs[1]["text"], paragraphs[2]["text"])) self.assertEqual(predictions, [None, "Syl", "Syl", "Syl", "Kaladin"]) def test_entity_relevance_counts_once_decays_and_stops_at_zero(self): scores = {"Syl": 0, "Kaladin": 0} update_entity_relevance(scores, "Syl saw Syl.") self.assertEqual(scores, {"Syl": 5, "Kaladin": 0}) update_entity_relevance(scores, "Kaladin arrived.") self.assertEqual(scores, {"Syl": 4, "Kaladin": 5}) for _ in range(10): update_entity_relevance(scores, "No named characters here.") self.assertEqual(scores, {"Syl": 0, "Kaladin": 0}) def test_unmarked_dialogue_alternates_by_relevance(self): text = ('Syl met Kaladin.\n\n' '"Hi," he said.\n\n' '"Within reason."\n\n' '"Whose reason?"') paragraphs = paragraphs_with_word_indexes(text) model = StubMaverick([[("Kaladin", 0), ("he", 0)]]) predictions, _ = evaluate_attributions( paragraphs, {}, set(), speaker_names=("Syl", "Kaladin"), maverick_model=model, ) self.assertEqual(predictions, [None, "Kaladin", "Syl", "Kaladin"]) def test_tied_relevance_leaves_unmarked_dialogue_other(self): text = 'Syl met Kaladin.\n\n"Hello."' paragraphs = paragraphs_with_word_indexes(text) predictions, _ = evaluate_attributions( paragraphs, {}, set(), speaker_names=("Syl", "Kaladin"), ) self.assertEqual(predictions, [None, "Other"]) def test_always_other_against_chapter_labels(self): for chapter_name, fixture_name, basename in FIXTURES: with self.subTest(chapter=chapter_name): fixture_dir = FIXTURE_ROOT / fixture_name text = (fixture_dir / f"{basename}.txt").read_bytes().decode("utf-8") paragraphs = paragraphs_with_word_indexes(text) labels = json.loads( (fixture_dir / f"{basename}.json").read_text(encoding="utf-8") ) italic_word_indexes = read_italic_word_indexes( fixture_dir / f"{basename}.rsvp", text ) expected_by_index = {} for chapter in labels["chapters"]: for annotation in chapter["paragraphs"]: index = annotation["paragraph_index"] - 1 self.assertGreaterEqual(index, 0) self.assertLess(index, len(paragraphs)) self.assertNotIn(index, expected_by_index) self.assertEqual( paragraphs[index]["text"].casefold(), annotation["text"].casefold(), ) expected_by_index[index] = annotation["entity"] predictions, stats = evaluate_attributions( paragraphs, expected_by_index, italic_word_indexes ) self.assertEqual(len(predictions), stats.total_paragraphs) self.assertTrue(all(entity in (None, "Other") for entity in predictions)) self.assertEqual( stats.total_paragraphs, stats.labeled_paragraphs + stats.unlabeled_paragraphs, ) self.assertEqual( stats.labeled_paragraphs, stats.correct_assignments + stats.incorrect_assignments + stats.missed_dialogue + stats.extra_dialogue + stats.correct_non_dialogue, ) self.assertEqual( stats.expected_dialogue, stats.correct_assignments + stats.incorrect_assignments + stats.missed_dialogue, ) self.assertEqual( stats.expected_non_dialogue, stats.correct_non_dialogue + stats.extra_dialogue, ) print(f"\nAttribution baseline: {chapter_name} — always Other") print(f"Paragraphs: {stats.total_paragraphs}") print(f"Explicitly labeled: {stats.labeled_paragraphs}") print(f"Unreviewed: {stats.unlabeled_paragraphs}") print(f"Expected dialogue: {stats.expected_dialogue}") print(f"Expected non-dialogue: {stats.expected_non_dialogue}") print(f"Correct assignments: {stats.correct_assignments}") print(f"Incorrect assignments: {stats.incorrect_assignments}") print(f"Missed dialogue: {stats.missed_dialogue}") print(f"Extra dialogue: {stats.extra_dialogue}") print(f"Correct non-dialogue: {stats.correct_non_dialogue}") print(f"Predicted dialogue on unreviewed: {stats.predicted_on_unlabeled}") @unittest.skipUnless(os.environ.get("RUN_MAVERICK_EVAL") == "1", "set RUN_MAVERICK_EVAL=1 to run the downloaded model") def test_maverick_against_chapter_labels(self): for chapter_name, fixture_name, basename in FIXTURES: with self.subTest(chapter=chapter_name): fixture_dir = FIXTURE_ROOT / fixture_name text = (fixture_dir / f"{basename}.txt").read_text(encoding="utf-8") paragraphs = paragraphs_with_word_indexes(text) labels = json.loads( (fixture_dir / f"{basename}.json").read_text(encoding="utf-8") ) italic_word_indexes = read_italic_word_indexes( fixture_dir / f"{basename}.rsvp", text ) expected_by_index = { annotation["paragraph_index"] - 1: annotation["entity"] for chapter in labels["chapters"] for annotation in chapter["paragraphs"] } predictions, stats = evaluate_attributions( paragraphs, expected_by_index, italic_word_indexes, speaker_names=tuple(labels["entities"].values()), ) if chapter_name == "Chapter 1": self.assertEqual(predictions[22], "Syl") self.assertEqual(predictions[24], "Syl") named = [ (index + 1, expected_by_index.get(index), prediction) for index, prediction in enumerate(predictions) if prediction not in (None, "Other") ] self.assertTrue(named) self.assertTrue(all(prediction in labels["entities"].values() for _, _, prediction in named)) print(f"\nAttribution rules: {chapter_name}") print(f"Correct assignments: {stats.correct_assignments}") print(f"Incorrect assignments: {stats.incorrect_assignments}") errors = [row for row in named if row[1] != row[2]] print(f"Named predictions: {len(named)}") print(f"Named errors (first 10): {errors[:10]}") if __name__ == "__main__": unittest.main()