417 lines
14 KiB
Python
417 lines
14 KiB
Python
#!/usr/bin/env python3
|
|
# Copyright 2026 Umbra Tools (umbratools.com)
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
"""Non-destructive chapter re-segmentation.
|
|
|
|
Re-runs the *current* Segmenter over selected chapters and merges the result
|
|
into the existing utterance table instead of wiping it. The goal is to apply
|
|
segmentation fixes without destroying data the user has already reviewed,
|
|
edited, or added:
|
|
|
|
* rows whose normalized text still matches a new segment keep their row id,
|
|
``speaker_name`` (including user-checked speakers), emotion/pace/instruct,
|
|
status, and ``user_edited`` flag;
|
|
* rows the user added/edited (``user_edited == 1`` or text not derivable from
|
|
the new segmentation) are always preserved even if the normalized text no
|
|
longer appears;
|
|
* narration rows replaced by corrected segmentation (leading-comma artifacts,
|
|
merged chapter titles, stale tense) are rewritten to the new text but stay
|
|
``narrator``;
|
|
* scene breaks stay ``scene_break`` and bind by order (old scenes are never
|
|
duplicated or dropped).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import contextlib
|
|
import json
|
|
import re
|
|
import sqlite3
|
|
import sys
|
|
import unicodedata
|
|
from difflib import SequenceMatcher
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
sys.path.insert(0, str(Path(__file__).parent.parent))
|
|
|
|
from umbraaudio.config import load_config
|
|
from umbraaudio.mapping.segmenter import Segmenter
|
|
|
|
|
|
def _norm(text: str) -> str:
|
|
"""Fold text for comparison: NFC, strip quotes, lowercase, collapse space."""
|
|
text = unicodedata.normalize("NFC", text or "")
|
|
text = text.strip(" \u201c\u201d\u2018\u2019\"'.,;:!?-\t\n")
|
|
for c in ("\u201c", "\u201d", "\u2018", "\u2019", '"', "'"):
|
|
text = text.replace(c, "")
|
|
return re.sub(r"\s+", " ", text.lower()).strip()
|
|
|
|
|
|
def _similar(a: str, b: str) -> float:
|
|
return SequenceMatcher(None, a, b).ratio()
|
|
|
|
|
|
def _row_fields(r: Any, rewritten: bool) -> dict:
|
|
return {
|
|
"text": r["text"],
|
|
"type": r["type"],
|
|
"speaker_name": r["speaker_name"],
|
|
"emotion": r["emotion"],
|
|
"pace": r["pace"],
|
|
"instruct": r["instruct"],
|
|
"status": r["status"],
|
|
"user_edited": r["user_edited"],
|
|
"attribution_confidence": r["attribution_confidence"],
|
|
"attribution_reasoning": r["attribution_reasoning"],
|
|
"context_before": r["context_before"],
|
|
"context_after": r["context_after"],
|
|
"rewritten": rewritten,
|
|
}
|
|
|
|
|
|
class _TargetPool:
|
|
"""Multiset of new segments keyed by normalized text."""
|
|
|
|
def __init__(self, new_segs: list[dict]) -> None:
|
|
self._groups: dict[str, list[dict]] = {}
|
|
for s in new_segs:
|
|
self._groups.setdefault(_norm(s.get("text", "")), []).append(s)
|
|
|
|
def exact(self, n: str) -> dict | None:
|
|
bucket = self._groups.get(n)
|
|
if not bucket:
|
|
return None
|
|
s = bucket.pop(0)
|
|
if not bucket:
|
|
del self._groups[n]
|
|
return s
|
|
|
|
def fuzzy_dialogue(self, n: str, threshold: float) -> dict | None:
|
|
best_key, best_sim, best = None, 0.0, None
|
|
for norm_key, cands in self._groups.items():
|
|
if "<<" in norm_key:
|
|
continue
|
|
sim = _similar(n, norm_key)
|
|
if sim > best_sim:
|
|
best_sim, best_key, best = sim, norm_key, cands[0]
|
|
if best is None or best_sim < threshold:
|
|
return None
|
|
bucket = self._groups.pop(best_key)
|
|
if len(bucket) > 1:
|
|
self._groups[best_key] = bucket[1:]
|
|
return best
|
|
|
|
def narration_rewrite(self, n: str) -> dict | None:
|
|
"""Match narration that got merged/split (title lines, comma fixes)."""
|
|
for norm_key, _cands in list(self._groups.items()):
|
|
if "<<" in norm_key:
|
|
continue
|
|
if n in norm_key or norm_key in n:
|
|
bucket = self._groups.pop(norm_key)
|
|
if len(bucket) > 1:
|
|
self._groups[norm_key] = bucket[1:]
|
|
return bucket[0]
|
|
return None
|
|
|
|
|
|
def _merge_chapter(
|
|
seg: Segmenter,
|
|
db_path: Path,
|
|
ch: Any,
|
|
character_names: list[str],
|
|
threshold: float,
|
|
dry_run: bool = False,
|
|
) -> tuple[int, int, int, list[int]]:
|
|
"""Rewrite one chapter's utterances in place. Returns (kept,new,rewritten,dropped)."""
|
|
ch_id = ch["id"]
|
|
title = ch["title"] or f"Chapter {ch['chapter_index']}"
|
|
source = ch["source_text"] or ""
|
|
new_segs = seg.segment(source, ch_id, character_names, chapter_title=title)
|
|
|
|
db = sqlite3.connect(str(db_path))
|
|
db.row_factory = sqlite3.Row
|
|
old_rows = db.execute(
|
|
"SELECT id, utterance_index, text, type, speaker_name, emotion, pace, "
|
|
"instruct, status, user_edited, attribution_confidence, "
|
|
"attribution_reasoning, context_before, context_after "
|
|
"FROM utterances WHERE chapter_id=? ORDER BY utterance_index",
|
|
(ch_id,),
|
|
).fetchall()
|
|
old_by_id: dict[int, Any] = {r["id"]: r for r in old_rows}
|
|
|
|
pool = _TargetPool(new_segs)
|
|
|
|
binds: dict[int, dict] = {}
|
|
for r in old_rows:
|
|
if r["type"] == "dialogue":
|
|
s = pool.exact(_norm(r["text"])) or pool.fuzzy_dialogue(_norm(r["text"]), threshold)
|
|
if s is not None:
|
|
binds[r["id"]] = s
|
|
for r in old_rows:
|
|
if r["type"] != "narration":
|
|
continue
|
|
s = pool.exact(_norm(r["text"])) or pool.narration_rewrite(_norm(r["text"]))
|
|
if s is not None:
|
|
binds[r["id"]] = s
|
|
|
|
# Scene breaks bind by order.
|
|
old_scenes = [r for r in old_rows if r["type"] == "scene_break"]
|
|
new_scenes = [s for s in new_segs if s.get("type") == "scene_break"]
|
|
scene_bind: dict[int, dict | None] = {
|
|
r["id"]: (new_scenes[i] if i < len(new_scenes) else None)
|
|
for i, r in enumerate(old_scenes)
|
|
}
|
|
|
|
assembly: list[dict] = []
|
|
consumed_old: set[int] = set()
|
|
|
|
def emit_row(rid: int, s: dict | None, scene: bool = False) -> None:
|
|
if rid in consumed_old:
|
|
return
|
|
r = old_by_id[rid]
|
|
fields = _row_fields(r, rewritten=False)
|
|
fields["old_index"] = r["utterance_index"]
|
|
if scene:
|
|
fields["text"] = "<<<>>>"
|
|
fields["type"] = "scene_break"
|
|
fields["speaker_name"] = "scene_break"
|
|
elif s is not None:
|
|
n_txt, n_tgt = _norm(r["text"]), _norm(s.get("text", ""))
|
|
fields["rewritten"] = bool(n_txt and n_tgt and n_txt != n_tgt)
|
|
if s.get("type") == "narration" and not r["user_edited"]:
|
|
fields["speaker_name"] = "narrator"
|
|
fields["text"] = s.get("text", r["text"])
|
|
fields["type"] = r["type"]
|
|
fields["new"] = False
|
|
fields["old_id"] = rid
|
|
assembly.append(fields)
|
|
consumed_old.add(rid)
|
|
|
|
for s in new_segs:
|
|
if s.get("type") == "scene_break":
|
|
rid = next((oid for oid, tgt in scene_bind.items() if tgt is s), None)
|
|
if rid is not None:
|
|
emit_row(rid, None, scene=True)
|
|
else:
|
|
assembly.append(
|
|
{
|
|
"text": "<<<>>>",
|
|
"type": "scene_break",
|
|
"speaker_name": "scene_break",
|
|
"new": True,
|
|
"old_id": None,
|
|
"rewritten": False,
|
|
}
|
|
)
|
|
continue
|
|
rid = next(
|
|
(oid for oid, tgt in binds.items() if tgt is s and oid not in consumed_old),
|
|
None,
|
|
)
|
|
if rid is not None:
|
|
emit_row(rid, s)
|
|
else:
|
|
assembly.append(
|
|
{
|
|
"text": s.get("text", ""),
|
|
"type": s.get("type", "narration"),
|
|
"speaker_name": s.get("character_name") or "narrator",
|
|
"new": True,
|
|
"old_id": None,
|
|
"rewritten": False,
|
|
}
|
|
)
|
|
|
|
for rid in sorted(scene_bind):
|
|
if scene_bind[rid] is None:
|
|
emit_row(rid, None, scene=True)
|
|
for rid in sorted(old_by_id):
|
|
if rid in consumed_old:
|
|
continue
|
|
r = old_by_id[rid]
|
|
n = _norm(r["text"])
|
|
derivable = any(
|
|
_similar(n, _norm(s.get("text", ""))) >= threshold for s in new_segs
|
|
)
|
|
if r["user_edited"] or not n or not derivable:
|
|
fields = _row_fields(r, rewritten=False)
|
|
fields["new"] = False
|
|
fields["old_id"] = rid
|
|
fields["old_index"] = r["utterance_index"]
|
|
assembly.append(fields)
|
|
consumed_old.add(rid)
|
|
|
|
dropped = [rid for rid in old_by_id if rid not in consumed_old]
|
|
|
|
kept = sum(1 for a in assembly if a["old_id"] is not None)
|
|
newly = sum(1 for a in assembly if a["new"])
|
|
rewritten = sum(1 for a in assembly if a["old_id"] is not None and a["rewritten"])
|
|
|
|
if dry_run:
|
|
db.close()
|
|
return kept, newly, rewritten, dropped
|
|
|
|
if dropped:
|
|
ph = ",".join("?" * len(dropped))
|
|
db.execute(f"DELETE FROM utterances WHERE id IN ({ph})", dropped)
|
|
for new_idx, a in enumerate(assembly):
|
|
if a["old_id"] is not None:
|
|
db.execute(
|
|
"UPDATE utterances SET utterance_index=?, text=?, speaker_name=?, "
|
|
"emotion=?, pace=?, instruct=?, status=? WHERE id=?",
|
|
(
|
|
new_idx,
|
|
a["text"],
|
|
a["speaker_name"],
|
|
a["emotion"],
|
|
a["pace"],
|
|
a["instruct"],
|
|
a["status"],
|
|
a["old_id"],
|
|
),
|
|
)
|
|
else:
|
|
db.execute(
|
|
"INSERT INTO utterances (chapter_id, utterance_index, text, type, "
|
|
"speaker_name, user_edited) VALUES (?, ?, ?, ?, ?, 0)",
|
|
(ch_id, new_idx, a["text"], a["type"], a["speaker_name"]),
|
|
)
|
|
db.commit()
|
|
db.close()
|
|
|
|
_remove_dropped_audio(db_path, ch_id, dropped, old_by_id)
|
|
_rename_chapter_audio(db_path, ch_id, assembly)
|
|
return kept, newly, rewritten, dropped
|
|
|
|
|
|
def _rename_chapter_audio(db_path: Path, chapter_id: int, assembly: list[dict]) -> None:
|
|
"""Rename utt_*.wav (+ sidecar .wav.txt) so audio follows the new indices."""
|
|
audio_dir = db_path.parent / "chapter_audio" / f"ch{chapter_id}"
|
|
if not audio_dir.exists():
|
|
return
|
|
index_map: dict[int, int] = {}
|
|
for new_idx, a in enumerate(assembly):
|
|
if a["old_id"] is not None and a.get("old_index") is not None:
|
|
index_map.setdefault(a["old_index"], new_idx)
|
|
if not index_map:
|
|
return
|
|
staged: list[tuple[Path, int]] = []
|
|
for file in sorted(audio_dir.glob("utt_*.wav")):
|
|
try:
|
|
old = int(file.stem.split("_")[1])
|
|
except (IndexError, ValueError):
|
|
continue
|
|
new = index_map.get(old)
|
|
if new is not None and new != old:
|
|
staged.append((file, new))
|
|
if not staged:
|
|
return
|
|
tmp_names: list[tuple[Path, int, Path]] = []
|
|
for file, new in staged:
|
|
tmp = audio_dir / f"._move_{file.stem}.wav.tmp"
|
|
file.rename(tmp)
|
|
tmp_names.append((tmp, new, file))
|
|
for tmp, new, orig in tmp_names:
|
|
tmp.rename(audio_dir / f"utt_{new:04d}.wav")
|
|
txt = audio_dir / f"{orig.stem}.wav.txt"
|
|
if txt.exists():
|
|
txt.rename(audio_dir / f"utt_{new:04d}.wav.txt")
|
|
|
|
|
|
def _remove_dropped_audio(
|
|
db_path: Path, chapter_id: int, dropped: list[int], old_by_id: dict
|
|
) -> None:
|
|
"""Delete wav/txt files whose utterances no longer exist to avoid stale audio."""
|
|
if not dropped:
|
|
return
|
|
audio_dir = db_path.parent / "chapter_audio" / f"ch{chapter_id}"
|
|
if not audio_dir.exists():
|
|
return
|
|
for rid in dropped:
|
|
r = old_by_id.get(rid)
|
|
if r is None:
|
|
continue
|
|
idx = r["utterance_index"]
|
|
for suffix in (".wav", ".wav.txt", ".mp3"):
|
|
f = audio_dir / f"utt_{idx:04d}{suffix}"
|
|
with contextlib.suppress(OSError):
|
|
f.unlink()
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--project", type=str, default=None)
|
|
parser.add_argument(
|
|
"--chapters",
|
|
type=str,
|
|
default=None,
|
|
help="Comma-separated chapter_index list, e.g. '6,7,8'",
|
|
)
|
|
parser.add_argument("--dry-run", action="store_true", help="Report changes without writing")
|
|
parser.add_argument(
|
|
"--match-threshold",
|
|
type=float,
|
|
default=0.82,
|
|
help="Similarity ratio required to bind a row (default 0.82)",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
cfgs = load_config(args.project)
|
|
db_path = cfgs.db_path
|
|
if not db_path.exists():
|
|
print(f"Database not found: {db_path}")
|
|
sys.exit(1)
|
|
|
|
speakers_path = cfgs.data_dir / "speakers.json"
|
|
character_names: list[str] = []
|
|
alias_map: dict[str, str] = {}
|
|
if speakers_path.exists():
|
|
with open(speakers_path) as f:
|
|
sdata = json.load(f)
|
|
for s in sdata.get("speakers", []):
|
|
pn = s.get("primary_name", "")
|
|
if pn:
|
|
character_names.append(pn)
|
|
for a in s.get("aliases", []):
|
|
alias_map[a] = pn
|
|
|
|
db = sqlite3.connect(str(db_path))
|
|
db.row_factory = sqlite3.Row
|
|
wanted = {int(x) for x in args.chapters.split(",") if x.strip()} if args.chapters else None
|
|
chapters = db.execute(
|
|
"SELECT id, title, chapter_index, source_text FROM chapters ORDER BY chapter_index"
|
|
).fetchall()
|
|
if wanted is not None:
|
|
chapters = [c for c in chapters if c["chapter_index"] in wanted]
|
|
db.close()
|
|
|
|
print(f"Segmenter re-run over {len(chapters)} chapters")
|
|
seg = Segmenter(alias_map=alias_map)
|
|
total_kept = total_new = total_dropped = total_rewritten = 0
|
|
|
|
for ch in chapters:
|
|
kept, newly, rewritten, dropped = _merge_chapter(
|
|
seg, db_path, ch, character_names, args.match_threshold, dry_run=args.dry_run
|
|
)
|
|
title = ch["title"] or ""
|
|
print(
|
|
f"ch{ch['chapter_index']:>2} {title[:26]:<26} "
|
|
f"kept={kept} new={newly} rewritten={rewritten} dropped={len(dropped)}"
|
|
)
|
|
total_kept += kept
|
|
total_new += newly
|
|
total_dropped += len(dropped)
|
|
total_rewritten += rewritten
|
|
|
|
print(
|
|
f"\nTotals: kept={total_kept} new={total_new} "
|
|
f"rewritten={total_rewritten} dropped={total_dropped}"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|