290 lines
12 KiB
Python
290 lines
12 KiB
Python
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()
|