146 lines
4.7 KiB
Python
146 lines
4.7 KiB
Python
#!/usr/bin/env python3
|
|
# Copyright 2026 Umbra Tools (umbratools.com)
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
"""Re-segment all chapters using mechanical splitter + LLM attribution."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import sqlite3
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, str(Path(__file__).parent.parent))
|
|
|
|
from umbraaudio.analysis.llm_attributor import LLMAttributor
|
|
from umbraaudio.config import LLMConfig
|
|
from umbraaudio.mapping.segmenter import Segmenter
|
|
|
|
|
|
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-12b-qat", help="LLM model name")
|
|
parser.add_argument(
|
|
"--mechanical-only", action="store_true", help="Skip LLM attribution, use mechanical only"
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
from umbraaudio.config import load_config
|
|
|
|
cfg = load_config(args.project)
|
|
config = cfg.llm
|
|
|
|
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, source_text FROM chapters "
|
|
f"WHERE id IN ({placeholders}) ORDER BY chapter_index",
|
|
chapter_ids,
|
|
).fetchall()
|
|
else:
|
|
chapters = db.execute(
|
|
"SELECT id, title, chapter_index, source_text FROM chapters ORDER BY chapter_index"
|
|
).fetchall()
|
|
|
|
print(f"Found {len(chapters)} chapters to segment")
|
|
print(f"Model: {args.model}")
|
|
print(f"LLM attribution: {'disabled' if args.mechanical_only else 'enabled'}")
|
|
|
|
# Load character names and alias map from speakers.json
|
|
speakers_path = cfg.data_dir / "speakers.json"
|
|
character_names = []
|
|
alias_map = {}
|
|
if speakers_path.exists():
|
|
with open(speakers_path) as f:
|
|
data = json.load(f)
|
|
for s in data.get("speakers", []):
|
|
primary = s.get("primary_name", "")
|
|
if not primary:
|
|
continue
|
|
character_names.append(primary)
|
|
for alias in s.get("aliases", []):
|
|
alias_map[alias] = primary
|
|
|
|
print(f"Loaded {len(character_names)} character names, {len(alias_map)} aliases")
|
|
|
|
seg = Segmenter(alias_map=alias_map)
|
|
attributor = LLMAttributor(config, model=args.model, data_dir=cfg.data_dir) if not args.mechanical_only else None
|
|
|
|
total_utterances = 0
|
|
|
|
for ch in chapters:
|
|
ch_id = ch["id"]
|
|
title = ch["title"] or f"Chapter {ch['chapter_index']}"
|
|
source_text = ch["source_text"]
|
|
|
|
old_count = db.execute(
|
|
"SELECT COUNT(*) as cnt FROM utterances WHERE chapter_id = ?", (ch_id,)
|
|
).fetchone()["cnt"]
|
|
|
|
# Phase 1: Mechanical segmentation
|
|
mechanical = seg.segment(source_text, ch_id, character_names, chapter_title=title)
|
|
|
|
utterances = []
|
|
for i, u in enumerate(mechanical):
|
|
speaker = u.get("character_name", "narrator")
|
|
if speaker is None:
|
|
speaker = "narrator"
|
|
|
|
is_scene_break = u["text"].strip() in ("<<<>>>",)
|
|
|
|
utterances.append(
|
|
{
|
|
"utterance_index": i,
|
|
"text": u["text"],
|
|
"type": u.get("type", "narration"),
|
|
"speaker_name": "scene_break" if is_scene_break else speaker,
|
|
}
|
|
)
|
|
|
|
# Phase 2: LLM attribution (if enabled)
|
|
if attributor and not args.mechanical_only:
|
|
utterances = attributor.attribute_speakers(utterances, title, ch_id)
|
|
|
|
# Write to database
|
|
db.execute("DELETE FROM utterances WHERE chapter_id = ?", (ch_id,))
|
|
|
|
for u in utterances:
|
|
db.execute(
|
|
"INSERT INTO utterances "
|
|
"(chapter_id, utterance_index, text, type, speaker_name) "
|
|
"VALUES (?, ?, ?, ?, ?)",
|
|
(ch_id, u["utterance_index"], u["text"], u["type"], u["speaker_name"]),
|
|
)
|
|
|
|
db.commit()
|
|
total_utterances += len(utterances)
|
|
|
|
print(f" {title}: {old_count} -> {len(utterances)} utterances")
|
|
|
|
db.close()
|
|
|
|
print(f"\nTotal: {total_utterances} utterances across {len(chapters)} chapters")
|
|
print("Done!")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|