322 lines
12 KiB
Python
322 lines
12 KiB
Python
#!/usr/bin/env python3
|
||
|
||
"""Interactively label dialogue paragraphs in converted chapter text.
|
||
|
||
Usage:
|
||
python3 book-processing/label_dialogue.py Chapter.txt
|
||
python3 book-processing/label_dialogue.py Chapter1.txt Chapter2.txt \
|
||
--output dialogue-labels.json
|
||
|
||
The script uses quotation marks in the text and italic word flags in matching
|
||
RSVP sidecars to find candidate paragraphs. It records labels by chapter and
|
||
paragraph index, with entity names as the stable output identity.
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import re
|
||
import struct
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
|
||
SIDECAR_HEADER = struct.Struct("<4sBBQI")
|
||
SIDECAR_MAGIC = b"RSVP"
|
||
SIDECAR_VERSIONS = {1, 2, 3}
|
||
STYLE_ITALIC = 0x01
|
||
STYLE_BOLD = 0x02
|
||
TIMING_MODIFIERS = {0, 2, 3}
|
||
LABEL_FORMAT = "dialogue-attribution-labels-v1"
|
||
|
||
# Single quotes are normalized to apostrophes by the converter. Require a
|
||
# pair that is not attached to word characters to avoid treating contractions
|
||
# such as don't as dialogue.
|
||
PAIRED_SINGLE_QUOTES = re.compile(r"(?<!\w)'[^'\n]+'(?!\w)")
|
||
WORD_TOKEN = re.compile(r"\b[\w’'-]+\b", re.UNICODE)
|
||
|
||
|
||
class SidecarError(ValueError):
|
||
"""A malformed, unsupported, or stale RSVP sidecar."""
|
||
|
||
|
||
def source_fingerprint(text):
|
||
"""Return the converter's stable 64-bit FNV-1a fingerprint."""
|
||
|
||
value = 0xCBF29CE484222325
|
||
for byte in text.encode("utf-8"):
|
||
value ^= byte
|
||
value = (value * 0x100000001B3) & 0xFFFFFFFFFFFFFFFF
|
||
return value
|
||
|
||
|
||
def read_varint(data, offset):
|
||
"""Read an unsigned LEB128 integer, returning its value and new offset."""
|
||
|
||
value = 0
|
||
shift = 0
|
||
while offset < len(data) and shift <= 28:
|
||
byte = data[offset]
|
||
offset += 1
|
||
if shift == 28 and byte > 0x0F:
|
||
raise SidecarError("sidecar word-index varint exceeds 32 bits")
|
||
value |= (byte & 0x7F) << shift
|
||
if not byte & 0x80:
|
||
return value, offset
|
||
shift += 7
|
||
raise SidecarError("invalid or oversized sidecar word-index varint")
|
||
|
||
|
||
def read_italic_word_indexes(sidecar_path, text):
|
||
"""Read italic word indexes after validating the sidecar against text."""
|
||
|
||
data = sidecar_path.read_bytes()
|
||
if len(data) < SIDECAR_HEADER.size:
|
||
raise SidecarError("sidecar is shorter than its header")
|
||
|
||
magic, version, reserved, fingerprint, record_count = SIDECAR_HEADER.unpack_from(data)
|
||
if magic != SIDECAR_MAGIC:
|
||
raise SidecarError("sidecar magic is not RSVP")
|
||
if version not in SIDECAR_VERSIONS:
|
||
raise SidecarError(f"unsupported sidecar version {version}")
|
||
if reserved != 0:
|
||
raise SidecarError("sidecar reserved header byte is not zero")
|
||
if fingerprint != source_fingerprint(text):
|
||
raise SidecarError("sidecar fingerprint does not match the chapter text")
|
||
|
||
word_count = sum(1 for _ in re.finditer(r"\S+", text))
|
||
if record_count > word_count:
|
||
raise SidecarError("sidecar has more records than chapter words")
|
||
|
||
italic_indexes = set()
|
||
offset = SIDECAR_HEADER.size
|
||
previous_index = 0
|
||
for record_number in range(record_count):
|
||
delta, offset = read_varint(data, offset)
|
||
record_bytes = 3 if version == 3 else 2
|
||
if offset + record_bytes > len(data):
|
||
raise SidecarError("sidecar ends inside a style record")
|
||
|
||
word_index = delta if record_number == 0 else previous_index + delta
|
||
flags, timing_modifier = data[offset], data[offset + 1]
|
||
color_flag = data[offset + 2] if version == 3 else 0
|
||
offset += record_bytes
|
||
|
||
if record_number and delta == 0:
|
||
raise SidecarError("sidecar word indexes are not strictly increasing")
|
||
if word_index >= word_count:
|
||
raise SidecarError("sidecar word index exceeds chapter word count")
|
||
if flags & ~(STYLE_ITALIC | STYLE_BOLD):
|
||
raise SidecarError("sidecar contains unknown style flags")
|
||
if timing_modifier not in TIMING_MODIFIERS:
|
||
raise SidecarError("sidecar contains an invalid timing modifier")
|
||
if version == 1 and timing_modifier != 0:
|
||
raise SidecarError("version 1 sidecar has a timing modifier")
|
||
if color_flag not in (0, 1):
|
||
raise SidecarError("sidecar contains an invalid word color flag")
|
||
if flags == 0 and timing_modifier == 0 and color_flag == 0:
|
||
raise SidecarError("sidecar contains an empty style record")
|
||
if flags & STYLE_ITALIC:
|
||
italic_indexes.add(word_index)
|
||
previous_index = word_index
|
||
|
||
if offset != len(data):
|
||
raise SidecarError("sidecar has trailing bytes")
|
||
return italic_indexes
|
||
|
||
|
||
def paragraphs_with_word_indexes(text):
|
||
"""Split blank-line-delimited paragraphs and report each word start index."""
|
||
|
||
paragraphs = []
|
||
word_start = 0
|
||
for raw_paragraph in re.split(r"\r?\n[ \t]*\r?\n+", text):
|
||
paragraph = raw_paragraph.strip()
|
||
if not paragraph:
|
||
continue
|
||
word_count = sum(1 for _ in re.finditer(r"\S+", paragraph))
|
||
paragraphs.append({
|
||
"index": len(paragraphs),
|
||
"text": paragraph,
|
||
"word_start": word_start,
|
||
"word_count": word_count,
|
||
})
|
||
word_start += word_count
|
||
return paragraphs
|
||
|
||
|
||
def paragraph_has_quotes(text):
|
||
"""Return whether a paragraph contains a likely dialogue quote marker."""
|
||
|
||
return '"' in text or bool(PAIRED_SINGLE_QUOTES.search(text))
|
||
|
||
|
||
def paragraph_has_all_caps_word(text):
|
||
"""Return whether a paragraph contains a word with 2+ uppercase letters."""
|
||
|
||
for token in WORD_TOKEN.findall(text):
|
||
letters = "".join(character for character in token if character.isalpha())
|
||
if len(letters) >= 2 and letters.isupper():
|
||
return True
|
||
return False
|
||
|
||
|
||
def is_candidate(paragraph, italic_indexes):
|
||
if paragraph_has_quotes(paragraph["text"]) or paragraph_has_all_caps_word(
|
||
paragraph["text"]
|
||
):
|
||
return True
|
||
start = paragraph["word_start"]
|
||
end = start + paragraph["word_count"]
|
||
return any(start <= index < end for index in italic_indexes)
|
||
|
||
|
||
def render_context(paragraphs, target_index, radius):
|
||
start = max(0, target_index - radius)
|
||
end = min(len(paragraphs), target_index + radius + 1)
|
||
for index in range(start, end):
|
||
label = "DIALOGUE" if index == target_index else "context"
|
||
print(f"[{label} paragraph {index + 1}]")
|
||
print(paragraphs[index]["text"])
|
||
print()
|
||
|
||
|
||
def save_labels(output_path, entities, chapters):
|
||
payload = {
|
||
"format": LABEL_FORMAT,
|
||
"entities": {str(number): name for number, name in sorted(entities.items())},
|
||
"chapters": chapters,
|
||
}
|
||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
temporary_path = output_path.with_name(output_path.name + ".tmp")
|
||
temporary_path.write_text(
|
||
json.dumps(payload, ensure_ascii=False, indent=2) + "\n",
|
||
encoding="utf-8",
|
||
)
|
||
temporary_path.replace(output_path)
|
||
|
||
|
||
def label_candidate(paragraphs, paragraph_index, entities):
|
||
radius = 1
|
||
while True:
|
||
render_context(paragraphs, paragraph_index, radius)
|
||
if entities:
|
||
print("Known entities: " + ", ".join(
|
||
f"{number}={name}" for number, name in sorted(entities.items())
|
||
))
|
||
print("Enter 0 for more context, 1-99 for a known entity, a name, or empty/null to skip.")
|
||
answer = input("Attribution: ").strip()
|
||
|
||
if answer == "0":
|
||
if radius >= len(paragraphs):
|
||
print("No additional chapter context is available.\n")
|
||
radius += 1
|
||
continue
|
||
if not answer or answer.casefold() == "null":
|
||
return None
|
||
if answer.isdecimal():
|
||
number = int(answer)
|
||
if 1 <= number <= 99 and number in entities:
|
||
return entities[number]
|
||
print("That number is not a previously assigned entity label.\n")
|
||
continue
|
||
|
||
known_number = next(
|
||
(number for number, name in entities.items()
|
||
if name.casefold() == answer.casefold()),
|
||
None,
|
||
)
|
||
if known_number is not None:
|
||
print(f"Using {known_number}={entities[known_number]}.\n")
|
||
return entities[known_number]
|
||
|
||
free_number = next((number for number in range(1, 100) if number not in entities), None)
|
||
if free_number is None:
|
||
print("All entity labels 1-99 are in use. Choose an existing label or skip.\n")
|
||
continue
|
||
entities[free_number] = answer
|
||
print(f"Assigned {free_number}={answer}.\n")
|
||
return answer
|
||
|
||
|
||
def load_chapter(path):
|
||
# Preserve newline bytes because the sidecar fingerprint covers the exact
|
||
# UTF-8 chapter text written by the converter.
|
||
text = path.read_bytes().decode("utf-8")
|
||
sidecar_path = path.with_suffix(".rsvp")
|
||
if sidecar_path.exists():
|
||
italic_indexes = read_italic_word_indexes(sidecar_path, text)
|
||
else:
|
||
print(f"Warning: {sidecar_path} is missing; italics will not be detected.", file=sys.stderr)
|
||
italic_indexes = set()
|
||
paragraphs = paragraphs_with_word_indexes(text)
|
||
if paragraphs and paragraphs[-1]["word_start"] + paragraphs[-1]["word_count"] != sum(
|
||
1 for _ in re.finditer(r"\S+", text)
|
||
):
|
||
raise ValueError(f"paragraph splitting did not align with word indexes in {path}")
|
||
return text, paragraphs, italic_indexes
|
||
|
||
|
||
def parse_args(argv=None):
|
||
parser = argparse.ArgumentParser(
|
||
description="Label quoted or italicized paragraphs with dialogue entities."
|
||
)
|
||
parser.add_argument("chapters", nargs="+", type=Path, help="chapter .txt files")
|
||
parser.add_argument(
|
||
"-o", "--output", type=Path, default=Path("dialogue-labels.json"),
|
||
help="output JSON file (default: dialogue-labels.json)",
|
||
)
|
||
return parser.parse_args(argv)
|
||
|
||
|
||
def main(argv=None):
|
||
args = parse_args(argv)
|
||
chapters = []
|
||
entities = {}
|
||
|
||
try:
|
||
for chapter_path in args.chapters:
|
||
if chapter_path.suffix.lower() != ".txt":
|
||
raise ValueError(f"expected a .txt chapter file: {chapter_path}")
|
||
if not chapter_path.is_file():
|
||
raise ValueError(f"chapter file does not exist: {chapter_path}")
|
||
|
||
text, paragraphs, italic_indexes = load_chapter(chapter_path)
|
||
chapter_result = {
|
||
"file": str(chapter_path.resolve()),
|
||
"paragraphs": [],
|
||
}
|
||
chapters.append(chapter_result)
|
||
save_labels(args.output, entities, chapters)
|
||
candidate_count = sum(
|
||
is_candidate(paragraph, italic_indexes) for paragraph in paragraphs
|
||
)
|
||
print(f"\n{chapter_path}: {candidate_count} candidate dialogue paragraphs")
|
||
|
||
for paragraph in paragraphs:
|
||
if not is_candidate(paragraph, italic_indexes):
|
||
continue
|
||
entity = label_candidate(paragraphs, paragraph["index"], entities)
|
||
chapter_result["paragraphs"].append({
|
||
"paragraph_index": paragraph["index"] + 1,
|
||
"text": paragraph["text"],
|
||
"entity": entity,
|
||
})
|
||
save_labels(args.output, entities, chapters)
|
||
|
||
save_labels(args.output, entities, chapters)
|
||
|
||
except (OSError, UnicodeError, SidecarError, ValueError) as exc:
|
||
print(f"Error: {exc}", file=sys.stderr)
|
||
return 2
|
||
except (EOFError, KeyboardInterrupt):
|
||
save_labels(args.output, entities, chapters)
|
||
print(f"\nSaved labels so far to {args.output}", file=sys.stderr)
|
||
return 130
|
||
|
||
print(f"\nSaved paragraph labels to {args.output}")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|