UmbraAudio/scripts/reattribute_all.py

#!/usr/bin/env python3
"""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.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 _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

        alias_map = _build_alias_map(cfg.data_dir / "speakers.json")

        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)
            ch_elapsed = time.time() - ch_start

            accepted_skipped = 0
            for u in attributed:
                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=? WHERE id=?",
                    (
                        u.get("speaker_name", "Narrator"),
                        u.get("attribution_confidence", u.get("confidence", 0.5)),
                        u.get("attribution_reasoning", "") or "",
                        u["id"],
                    ),
                )
            if not args.dry_run:
                db.commit()

            if args.dry_run:
                _report_dry_run(attributed, 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()