198 lines
7.2 KiB
Python
198 lines
7.2 KiB
Python
# Copyright 2026 Umbra Tools (umbratools.com)
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
from __future__ import annotations
|
|
|
|
from umbraaudio.analysis.attribution_qa import (
|
|
check_attribution,
|
|
fix_high_confidence,
|
|
)
|
|
|
|
ALIAS_MAP = {
|
|
"hunter": "Deal Hunter",
|
|
"deal hunter": "Deal Hunter",
|
|
"kainda": "Kainda Walters Foster Murphy",
|
|
"kainda walters foster murphy": "Kainda Walters Foster Murphy",
|
|
"princess kainda": "Kainda Walters Foster Murphy",
|
|
"kabir": "Kabir",
|
|
"narrator": "Narrator",
|
|
}
|
|
|
|
|
|
def _utt(idx: int, text: str, speaker: str, type_: str = "narration") -> dict:
|
|
return {
|
|
"utterance_index": idx,
|
|
"text": text,
|
|
"type": type_,
|
|
"speaker_name": speaker,
|
|
}
|
|
|
|
|
|
class TestSpeechTagMismatch:
|
|
def test_correct_attribution_no_issue(self) -> None:
|
|
utts = [
|
|
_utt(0, "Hunter says,", "narrator"),
|
|
_utt(1, "We are ready to jump.", "Deal Hunter", "dialogue"),
|
|
]
|
|
result = check_attribution(utts, 6, ALIAS_MAP)
|
|
assert result.has("speech_tag_mismatch") is False
|
|
|
|
def test_wrong_attribution_flagged(self) -> None:
|
|
utts = [
|
|
_utt(0, "Hunter says,", "narrator"),
|
|
_utt(1, "We are ready to jump.", "Kabir", "dialogue"),
|
|
]
|
|
result = check_attribution(utts, 6, ALIAS_MAP)
|
|
issues = result.by_type("speech_tag_mismatch")
|
|
assert len(issues) == 1
|
|
assert issues[0].expected == "Deal Hunter"
|
|
assert issues[0].utterance_index == 1
|
|
assert issues[0].confidence >= 0.9
|
|
|
|
def test_alias_resolution(self) -> None:
|
|
utts = [
|
|
_utt(0, "Princess Kainda asks,", "narrator"),
|
|
_utt(1, "Where are we?", "Deal Hunter", "dialogue"),
|
|
]
|
|
result = check_attribution(utts, 6, ALIAS_MAP)
|
|
issues = result.by_type("speech_tag_mismatch")
|
|
assert issues[0].expected == "Kainda Walters Foster Murphy"
|
|
|
|
def test_tag_not_carried_across_unrelated_narration(self) -> None:
|
|
utts = [
|
|
_utt(0, "Hunter says,", "narrator"),
|
|
_utt(1, "The ship shudders.", "narrator"),
|
|
_utt(2, "We are ready to jump.", "Kabir", "dialogue"),
|
|
]
|
|
result = check_attribution(utts, 6, ALIAS_MAP)
|
|
assert result.has("speech_tag_mismatch") is False
|
|
|
|
def test_tag_not_carried_across_scene_break(self) -> None:
|
|
utts = [
|
|
_utt(0, "Hunter says,", "narrator"),
|
|
_utt(1, "<<<>>>", "scene_break"),
|
|
_utt(2, "We are ready to jump.", "Kabir", "dialogue"),
|
|
]
|
|
result = check_attribution(utts, 6, ALIAS_MAP)
|
|
assert result.has("speech_tag_mismatch") is False
|
|
|
|
|
|
class TestNarratorDialogue:
|
|
def test_dialogue_left_as_narrator_flagged(self) -> None:
|
|
utts = [
|
|
_utt(0, "She says,", "narrator"),
|
|
_utt(1, "Hello.", "narrator", "dialogue"),
|
|
]
|
|
result = check_attribution(utts, 6, ALIAS_MAP)
|
|
issues = result.by_type("narrator_dialogue")
|
|
assert len(issues) == 1
|
|
assert issues[0].utterance_index == 1
|
|
|
|
|
|
class TestRoleTag:
|
|
def test_role_tag_assigned_to_named_character_flagged(self) -> None:
|
|
utts = [
|
|
_utt(0, "The lead guard replies,", "narrator"),
|
|
_utt(1, "We have strict orders to keep you away.", "Kabir", "dialogue"),
|
|
]
|
|
result = check_attribution(utts, 6, ALIAS_MAP)
|
|
issues = result.by_type("role_tag_guessed")
|
|
assert len(issues) == 1
|
|
assert issues[0].utterance_index == 1
|
|
|
|
def test_role_tag_assigned_to_role_speaker_ok(self) -> None:
|
|
utts = [
|
|
_utt(0, "The lead guard replies,", "narrator"),
|
|
_utt(1, "We have strict orders to keep you away.", "Chief of the Guard", "dialogue"),
|
|
]
|
|
result = check_attribution(utts, 6, ALIAS_MAP)
|
|
assert result.has("role_tag_guessed") is False
|
|
|
|
def test_pronoun_tag_skipped(self) -> None:
|
|
utts = [
|
|
_utt(0, "He replies,", "narrator"),
|
|
_utt(1, "We have strict orders to keep you away.", "Chief of the Guard", "dialogue"),
|
|
]
|
|
result = check_attribution(utts, 6, ALIAS_MAP)
|
|
assert not result.issues
|
|
|
|
def test_tag_not_applied_to_shared_tag_block(self) -> None:
|
|
utts = [
|
|
_utt(0, "Kainda asks,", "narrator"),
|
|
_utt(
|
|
1, "Where is the ship that brought me?", "Kainda Walters Foster Murphy", "dialogue"
|
|
),
|
|
_utt(
|
|
2,
|
|
"The ship jumped away as soon as you were on your way.",
|
|
"Willamina Wake",
|
|
"dialogue",
|
|
),
|
|
]
|
|
result = check_attribution(utts, 6, ALIAS_MAP)
|
|
assert not result.issues
|
|
|
|
def test_voice_tag_assigned_to_ship_flagged(self) -> None:
|
|
utts = [
|
|
_utt(0, "The voice calls again,", "narrator"),
|
|
_utt(1, "Help me, please.", "Deal Hunter", "dialogue"),
|
|
]
|
|
result = check_attribution(utts, 6, ALIAS_MAP)
|
|
issues = result.by_type("role_tag_guessed")
|
|
assert len(issues) == 1
|
|
|
|
def test_unresolved_title_word_overlap_ok(self) -> None:
|
|
utts = [
|
|
_utt(0, "Wakam says,", "narrator"),
|
|
_utt(1, "You work a tough deal.", "Trader Wakam", "dialogue"),
|
|
]
|
|
result = check_attribution(utts, 6, ALIAS_MAP)
|
|
assert not result.issues
|
|
|
|
def test_unresolved_title_no_overlap_flagged_low_conf(self) -> None:
|
|
utts = [
|
|
_utt(0, "Wakam says,", "narrator"),
|
|
_utt(1, "You work a tough deal.", "Captain Weston", "dialogue"),
|
|
]
|
|
result = check_attribution(utts, 6, ALIAS_MAP)
|
|
issues = result.by_type("speech_tag_mismatch")
|
|
assert len(issues) == 1
|
|
assert issues[0].confidence < 0.9
|
|
|
|
def test_conjunction_not_absorbed_into_role(self) -> None:
|
|
utts = [
|
|
_utt(0, "A supervisor sees the robot and tells it to stop.", "narrator"),
|
|
_utt(1, "Why are you here?", "Human Staff Person/Messenger", "dialogue"),
|
|
]
|
|
result = check_attribution(utts, 6, ALIAS_MAP)
|
|
assert not result.issues
|
|
|
|
def test_plain_narration_not_flagged(self) -> None:
|
|
utts = [
|
|
_utt(0, "The ship powers down its engines.", "narrator"),
|
|
_utt(1, "Help me, please.", "Deal Hunter", "dialogue"),
|
|
]
|
|
result = check_attribution(utts, 6, ALIAS_MAP)
|
|
assert not result.issues
|
|
|
|
|
|
class TestFixHighConfidence:
|
|
def test_fixes_exact_name_mismatch_only(self) -> None:
|
|
utts = [
|
|
_utt(0, "Hunter says,", "narrator"),
|
|
_utt(1, "We are ready to jump.", "Kabir", "dialogue"),
|
|
_utt(2, "The ship shudders.", "narrator"),
|
|
_utt(3, "Can you take me on a tour?", "Kabir", "dialogue"),
|
|
]
|
|
fixed = fix_high_confidence(utts, ALIAS_MAP)
|
|
assert fixed[1]["speaker_name"] == "Deal Hunter"
|
|
assert fixed[3]["speaker_name"] == "Kabir"
|
|
|
|
def test_already_correct_unchanged(self) -> None:
|
|
utts = [
|
|
_utt(0, "Hunter says,", "narrator"),
|
|
_utt(1, "We are ready to jump.", "Deal Hunter", "dialogue"),
|
|
]
|
|
fixed = fix_high_confidence(utts, ALIAS_MAP)
|
|
assert fixed[1]["speaker_name"] == "Deal Hunter"
|