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 sqlite3
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

from umbraaudio.analysis.llm_attributor import LLMAttributor
from umbraaudio.config import AppConfig, LLMConfig
from umbraaudio.voices.model_coordinator import ModelCoordinator, PipelineStage


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(
        "--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",
    )
    args = parser.parse_args()

    from umbraaudio.config import load_config

    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)

    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

    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()

    print(f"Found {len(chapters)} chapters to attribute")
    print(f"Model: {args.model}")

    attributor = LLMAttributor(config, model=args.model, data_dir=cfg.data_dir)
    total_dialogue = 0
    total_narrator_left = 0

    for ch in chapters:
        ch_id = ch["id"]
        title = ch["title"] or f"Chapter {ch['chapter_index']}"

        utts = db.execute(
            "SELECT id, utterance_index, text, type, speaker_name "
            "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"
        )

        attributed = attributor.attribute_speakers(utt_list, title, ch_id)
        time.sleep(0.5)

        for u in attributed:
            db.execute(
                "UPDATE utterances SET speaker_name=? WHERE id=?",
                (u.get("speaker_name", "Narrator"), u["id"]),
            )
        db.commit()

        narrator_after = sum(
            1
            for u in attributed
            if u.get("type") == "dialogue"
            and (u.get("speaker_name") or "narrator").lower() == "narrator"
        )
        total_dialogue += dialogue_before
        total_narrator_left += narrator_after
        print(
            f"  {title} (id={ch_id}): dialogue {dialogue_before}, "
            f"narrator-left {narrator_before} -> {narrator_after}"
        )

    db.close()

    coordinator.switch_to(PipelineStage.IDLE)

    print(
        f"\nTotal: {total_dialogue} dialogue utterances, "
        f"{total_narrator_left} still attributed to narrator"
    )
    print("Done!")


if __name__ == "__main__":
    main()