320 lines
11 KiB
Python
320 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
# Copyright 2026 Umbra Tools (umbratools.com)
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
"""Re-run LLM speaker attribution on existing utterances without re-segmenting.
|
|
|
|
Preserves utterance IDs and mechanical segmentation; only updates speaker_name.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import sqlite3
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, str(Path(__file__).parent.parent))
|
|
|
|
from umbraaudio.analysis.attribution_qa import (
|
|
check_attribution,
|
|
verify_conversation_flow,
|
|
)
|
|
from umbraaudio.analysis.clarity import assess_utterance_clarity
|
|
from umbraaudio.analysis.llm_attributor import LLMAttributor
|
|
from umbraaudio.config import load_config
|
|
from umbraaudio.voice_maintenance import (
|
|
clear_voice_maintenance,
|
|
set_voice_maintenance,
|
|
)
|
|
from umbraaudio.voices.model_coordinator import ModelCoordinator, PipelineStage
|
|
|
|
|
|
def _ensure_utterance_columns(db: sqlite3.Connection) -> None:
|
|
"""Add clarity columns introduced after project DBs were first created."""
|
|
cols = {r[1] for r in db.execute("PRAGMA table_info(utterances)")}
|
|
for col, typedef in (
|
|
("attribution_evidence_kind", "TEXT"),
|
|
("attribution_tag_name", "TEXT"),
|
|
("attribution_uncertain", "INTEGER NOT NULL DEFAULT 0"),
|
|
("needs_review", "INTEGER NOT NULL DEFAULT 0"),
|
|
("review_reason", "TEXT"),
|
|
):
|
|
if col not in cols:
|
|
db.execute(f"ALTER TABLE utterances ADD COLUMN {col} {typedef}")
|
|
db.commit()
|
|
|
|
|
|
def _build_alias_map(speakers_path: Path) -> dict[str, str]:
|
|
if not speakers_path.exists():
|
|
return {}
|
|
with open(speakers_path) as f:
|
|
data = json.load(f)
|
|
alias_map: dict[str, str] = {}
|
|
for s in data.get("speakers", []):
|
|
primary = s.get("primary_name", "")
|
|
if not primary:
|
|
continue
|
|
alias_map[primary.lower()] = primary
|
|
for alias in s.get("aliases", []):
|
|
if alias:
|
|
alias_map[alias.lower()] = primary
|
|
return alias_map
|
|
|
|
|
|
def _qa_issue_count(utt_list: list[dict], chapter_id: int, alias_map: dict) -> int:
|
|
qa = check_attribution(utt_list, chapter_id, alias_map)
|
|
flow = verify_conversation_flow(utt_list, chapter_id, alias_map)
|
|
return len(qa.issues) + len(flow.issues)
|
|
|
|
|
|
def _report_dry_run(attributed: list[dict], original: list[dict]) -> None:
|
|
"""Print speaker changes the run would make (no DB writes)."""
|
|
orig_by_id = {u["id"]: u for u in original}
|
|
changed = 0
|
|
shown = 0
|
|
for u in attributed:
|
|
o = orig_by_id.get(u["id"])
|
|
if o is None:
|
|
continue
|
|
before = o.get("speaker_name") or "narrator"
|
|
after = u.get("speaker_name") or "narrator"
|
|
if (before or "").lower() == (after or "").lower():
|
|
continue
|
|
changed += 1
|
|
if shown < 5:
|
|
text = (u.get("text") or "")[:70]
|
|
print(f" utt {u.get('utterance_index')}: {before} -> {after} {text!r}")
|
|
shown += 1
|
|
print(f" dry-run: {changed} speaker change(s) (sample above)")
|
|
|
|
|
|
def main():
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--project",
|
|
type=str,
|
|
default=None,
|
|
help="Project name or path",
|
|
)
|
|
parser.add_argument(
|
|
"--chapters",
|
|
type=str,
|
|
default=None,
|
|
help="Comma-separated chapter IDs (e.g. '5,6,7')",
|
|
)
|
|
parser.add_argument(
|
|
"--min-id",
|
|
type=int,
|
|
default=None,
|
|
help="Only chapters with id >= this value (e.g. 11 = 'chapters 11 to the end')",
|
|
)
|
|
parser.add_argument(
|
|
"--max-id",
|
|
type=int,
|
|
default=None,
|
|
help="Only chapters with id <= this value",
|
|
)
|
|
parser.add_argument(
|
|
"--model",
|
|
type=str,
|
|
default="google/gemma-4-e4b",
|
|
help="LLM model name",
|
|
)
|
|
parser.add_argument(
|
|
"--sleep",
|
|
type=float,
|
|
default=1.0,
|
|
help="Seconds to sleep between LLM batches",
|
|
)
|
|
parser.add_argument(
|
|
"--no-verify",
|
|
action="store_true",
|
|
help="Skip the multi-pass verification round (pass 3)",
|
|
)
|
|
parser.add_argument(
|
|
"--verify-rounds",
|
|
type=int,
|
|
default=2,
|
|
help="Max verification rounds for flagged lines (default 2)",
|
|
)
|
|
parser.add_argument(
|
|
"--dry-run",
|
|
action="store_true",
|
|
help="Attribute in memory and report vs DB, but write NOTHING",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
cfg = load_config(args.project)
|
|
config = cfg.llm
|
|
|
|
coordinator = ModelCoordinator(cfg)
|
|
if not coordinator.switch_to(
|
|
PipelineStage.CHAPTER_ANALYSIS, lm_model=args.model
|
|
).get("ok"):
|
|
print("Failed to switch models to CHAPTER_ANALYSIS stage")
|
|
sys.exit(1)
|
|
|
|
set_voice_maintenance(
|
|
"LLM speaker re-attribution is running — Qwen3-TTS is offline "
|
|
"for the analysis. Voice re-generation is paused until the run finishes."
|
|
)
|
|
|
|
try:
|
|
db_path = cfg.db_path
|
|
if not db_path.exists():
|
|
print(f"Database not found: {db_path}")
|
|
return
|
|
|
|
db = sqlite3.connect(str(db_path))
|
|
db.row_factory = sqlite3.Row
|
|
_ensure_utterance_columns(db)
|
|
|
|
alias_map = _build_alias_map(cfg.data_dir / "speakers.json")
|
|
profiles: list[dict] = []
|
|
speakers_path = cfg.data_dir / "speakers.json"
|
|
if speakers_path.exists():
|
|
with open(speakers_path) as f:
|
|
profiles = json.load(f).get("speakers", [])
|
|
|
|
if args.chapters:
|
|
chapter_ids = [int(x.strip()) for x in args.chapters.split(",")]
|
|
placeholders = ",".join("?" * len(chapter_ids))
|
|
chapters = db.execute(
|
|
f"SELECT id, title, chapter_index FROM chapters "
|
|
f"WHERE id IN ({placeholders}) ORDER BY chapter_index",
|
|
chapter_ids,
|
|
).fetchall()
|
|
else:
|
|
chapters = db.execute(
|
|
"SELECT id, title, chapter_index FROM chapters ORDER BY chapter_index"
|
|
).fetchall()
|
|
|
|
if args.min_id is not None or args.max_id is not None:
|
|
chapters = [
|
|
ch
|
|
for ch in chapters
|
|
if (args.min_id is None or ch["id"] >= args.min_id)
|
|
and (args.max_id is None or ch["id"] <= args.max_id)
|
|
]
|
|
|
|
print(f"Found {len(chapters)} chapters to attribute")
|
|
print(f"Model: {args.model}")
|
|
print(f"Verify pass: {'ON' if not args.no_verify else 'OFF'} "
|
|
f"(up to {args.verify_rounds} round(s))")
|
|
if args.dry_run:
|
|
print("DRY RUN: no changes will be written")
|
|
|
|
attributor = LLMAttributor(config, model=args.model, data_dir=cfg.data_dir)
|
|
total_dialogue = 0
|
|
total_narrator_left = 0
|
|
total_qa_before = 0
|
|
total_qa_after = 0
|
|
|
|
for ch in chapters:
|
|
ch_id = ch["id"]
|
|
title = ch["title"] or f"Chapter {ch['chapter_index']}"
|
|
ch_start = time.time()
|
|
|
|
utts = db.execute(
|
|
"SELECT id, utterance_index, text, type, speaker_name, "
|
|
"qa_accepted, qa_accepted_sig "
|
|
"FROM utterances WHERE chapter_id=? ORDER BY utterance_index",
|
|
(ch_id,),
|
|
).fetchall()
|
|
if not utts:
|
|
print(f" {title}: no utterances, skipping")
|
|
continue
|
|
|
|
utt_list = [dict(u) for u in utts]
|
|
dialogue_before = sum(1 for u in utt_list if u.get("type") == "dialogue")
|
|
narrator_before = sum(
|
|
1
|
|
for u in utt_list
|
|
if u.get("type") == "dialogue"
|
|
and (u.get("speaker_name") or "narrator").lower() == "narrator"
|
|
)
|
|
qa_before = _qa_issue_count(utt_list, ch_id, alias_map)
|
|
|
|
attributed = attributor.attribute_speakers(
|
|
utt_list,
|
|
title,
|
|
ch_id,
|
|
verify=not args.no_verify,
|
|
verify_rounds=args.verify_rounds,
|
|
)
|
|
time.sleep(0.5)
|
|
|
|
clarity_aware = assess_utterance_clarity(
|
|
attributed, ch_id, alias_map, profiles
|
|
)
|
|
ch_elapsed = time.time() - ch_start
|
|
|
|
accepted_skipped = 0
|
|
for u in clarity_aware:
|
|
if u.get("qa_accepted"):
|
|
accepted_skipped += 1
|
|
continue
|
|
if args.dry_run:
|
|
continue
|
|
db.execute(
|
|
"UPDATE utterances SET speaker_name=?, attribution_confidence=?, "
|
|
"attribution_reasoning=?, attribution_evidence_kind=?, "
|
|
"attribution_tag_name=?, attribution_uncertain=?, "
|
|
"needs_review=?, review_reason=? WHERE id=?",
|
|
(
|
|
u.get("speaker_name", "Narrator"),
|
|
u.get("attribution_confidence", u.get("confidence", 0.5)),
|
|
u.get("attribution_reasoning", "") or "",
|
|
u.get("attribution_evidence_kind") or None,
|
|
u.get("attribution_tag_name") or None,
|
|
1 if u.get("attribution_uncertain") else 0,
|
|
1 if u.get("needs_review") else 0,
|
|
u.get("review_reason") or "",
|
|
u["id"],
|
|
),
|
|
)
|
|
if not args.dry_run:
|
|
db.commit()
|
|
|
|
if args.dry_run:
|
|
_report_dry_run(clarity_aware, utt_list)
|
|
|
|
narrator_after = sum(
|
|
1
|
|
for u in attributed
|
|
if u.get("type") == "dialogue"
|
|
and (u.get("speaker_name") or "narrator").lower() == "narrator"
|
|
)
|
|
qa_after = _qa_issue_count(attributed, ch_id, alias_map)
|
|
total_dialogue += dialogue_before
|
|
total_narrator_left += narrator_after
|
|
total_qa_before += qa_before
|
|
total_qa_after += qa_after
|
|
print(
|
|
f" {title} (id={ch_id}): dialogue {dialogue_before}, "
|
|
f"narrator-left {narrator_before} -> {narrator_after}, "
|
|
f"QA issues {qa_before} -> {qa_after}"
|
|
f" ({ch_elapsed:.0f}s)"
|
|
+ (f", accepted-skipped {accepted_skipped}" if accepted_skipped else "")
|
|
)
|
|
|
|
db.close()
|
|
|
|
print(
|
|
f"\nTotal: {total_dialogue} dialogue utterances, "
|
|
f"{total_narrator_left} still attributed to narrator, "
|
|
f"QA issues {total_qa_before} -> {total_qa_after}"
|
|
)
|
|
finally:
|
|
clear_voice_maintenance()
|
|
coordinator.switch_to(PipelineStage.IDLE)
|
|
|
|
print("Done!")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|