441 lines
17 KiB
Python
441 lines
17 KiB
Python
#!/usr/bin/env python3
|
||
"""exp15 — chunkers for Q1: a faithful zh source-sentence splitter + the §B1.5 LOGICAL DP chunker.
|
||
|
||
Q1 contrast (design §D2): A0 (greedy-1500, from the real backend) vs a LOGICAL-boundary arm sized to
|
||
A0's ACTUAL est_out distribution (±10% mean). The ONLY variable is cut ALIGNMENT (greedy is blind to
|
||
scene structure; logical aligns cuts to scene-shift boundaries and avoids mid-scene/dialogue cuts).
|
||
Size is held equal so Q1 does not confound boundary class with chunk size (that is Q2's domain).
|
||
|
||
Ports (backend/internal/pipeline/chunker.go):
|
||
splitSourceSentences / isOpenQuote / isCloseQuote / isSourceTerminator / isCJKTerminator /
|
||
isClosingBracket / isASCIISpace (chunker.go:196-357) — quote-depth-aware, lossless tiling.
|
||
Token math:
|
||
est_tokens = cjk + other//3 (EstimateTokens units — the greedy 1500 budget space)
|
||
est_out = 1.1978*cjk + 0.3852*other (fertility_calib.json — glm-4.6 ru output tokens; §B1.2)
|
||
|
||
The logical chunker implements §B1.4 boundary classes (B0 chapter/separator, B1 scene-shift lexicon,
|
||
B2 TextTiling cohesion valley, B3 paragraph, B4 sentence) and §B1.5 DP packing (Knuth-Plass shortest
|
||
path over candidate cuts, INTEGER-scaled costs, asymmetric size penalty). It does NOT peek at the
|
||
ground-truth seam offsets — scene-shifts are DISCOVERED (that is the Q1 test).
|
||
|
||
Self-review by execution: `python chunker.py --selftest` checks tiling losslessness + a size sweep.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import math
|
||
import sys
|
||
import unicodedata
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
|
||
import regex
|
||
|
||
_RE_CJK = regex.compile(r"[\p{Han}\p{Hiragana}\p{Katakana}\p{Hangul}]")
|
||
TARGET_CHUNK_TOKENS = 1500 # chunker.go:47 const
|
||
|
||
F_CJK, F_OTHER = 1.1978, 0.3852 # fertility_calib.json (glm-4.6 ru out tokens from src char classes)
|
||
|
||
# --- char-class token estimators -------------------------------------------------------------------
|
||
def _classes(text: str):
|
||
cjk = other = 0
|
||
for ch in text:
|
||
if _RE_CJK.match(ch):
|
||
cjk += 1
|
||
elif ch.isspace():
|
||
pass
|
||
else:
|
||
other += 1
|
||
return cjk, other
|
||
|
||
|
||
def est_tokens(text: str) -> int:
|
||
"""EstimateTokens units (source side): cjk + other//3, floored at 16 (chunker.go estTokensFrom)."""
|
||
cjk, other = _classes(text)
|
||
v = cjk + other // 3
|
||
return max(v, 16)
|
||
|
||
|
||
def est_out(text: str) -> float:
|
||
"""Predicted glm-4.6 ru OUTPUT tokens for a source slice (fertility model; §B1.2)."""
|
||
cjk, other = _classes(text)
|
||
return F_CJK * cjk + F_OTHER * other
|
||
|
||
|
||
# --- GREEDY chunker port (chunker.go:82-194 + render.go:285 NormalizeSource) -----------------------
|
||
def normalize_source(s: str) -> str:
|
||
"""render.go:285 — BOM strip, CRLF/CR->LF, valid-UTF8 (Python strs already are), NFC, TrimSpace."""
|
||
s = s.lstrip("")
|
||
s = s.replace("\r\n", "\n").replace("\r", "\n")
|
||
s = unicodedata.normalize("NFC", s)
|
||
return s.strip() # Go strings.TrimSpace strips Unicode whitespace incl. U+3000 (== Python str.strip())
|
||
|
||
|
||
def _class_counts(s: str):
|
||
cjk = other = 0
|
||
for ch in s:
|
||
if _RE_CJK.match(ch):
|
||
cjk += 1
|
||
elif ch.isspace():
|
||
pass
|
||
else:
|
||
other += 1
|
||
return cjk, other
|
||
|
||
|
||
def _est_tokens_from(cjk, other):
|
||
return max(cjk + other // 3, 16) # chunker.go estTokensFrom (integer // 3, floor 16)
|
||
|
||
|
||
def split_paragraphs(chapter: str):
|
||
return [p.strip() for p in chapter.split("\n\n") if p.strip()]
|
||
|
||
|
||
def _pack_sentences(paragraph: str):
|
||
segs = split_source_sentences(paragraph)
|
||
chunks, buf, bc, bo = [], [], 0, 0
|
||
def flush():
|
||
nonlocal buf, bc, bo
|
||
if buf:
|
||
chunks.append("".join(buf)); buf, bc, bo = [], 0, 0
|
||
for s in segs:
|
||
sc, so = _class_counts(s)
|
||
if buf and _est_tokens_from(bc + sc, bo + so) > TARGET_CHUNK_TOKENS:
|
||
flush()
|
||
buf.append(s); bc, bo = bc + sc, bo + so
|
||
flush()
|
||
return chunks
|
||
|
||
|
||
def greedy_chunks(raw_text: str):
|
||
"""Port of chunker.go SplitChunks for ONE chapter (S2' has no chapter markers). Returns chunk texts."""
|
||
chapter = normalize_source(raw_text)
|
||
paras = split_paragraphs(chapter)
|
||
chunks, buf, bc, bo = [], [], 0, 0
|
||
def flush():
|
||
nonlocal buf, bc, bo
|
||
if buf:
|
||
text = "\n\n".join(buf).strip()
|
||
if text:
|
||
chunks.append(text)
|
||
buf, bc, bo = [], 0, 0
|
||
for p in paras:
|
||
pc, po = _class_counts(p)
|
||
if _est_tokens_from(pc, po) > TARGET_CHUNK_TOKENS:
|
||
flush()
|
||
for sub in _pack_sentences(p):
|
||
t = sub.strip()
|
||
if t:
|
||
chunks.append(t)
|
||
continue
|
||
if buf and _est_tokens_from(bc + pc, bo + po) > TARGET_CHUNK_TOKENS:
|
||
flush()
|
||
buf.append(p); bc, bo = bc + pc, bo + po
|
||
flush()
|
||
return chunks
|
||
|
||
|
||
# --- source sentence splitter (port of chunker.go:196-357) -----------------------------------------
|
||
_OPEN_Q = set("「『（【《〈［｛“‘")
|
||
_CLOSE_Q = set("」』）】》〉］｝”’")
|
||
_TERM = set("。！？.!?…")
|
||
_CJK_TERM = set("。！？")
|
||
_CLOSE_BRK = set("」』】）)》〉］]｝}\"'”’")
|
||
_ASCII_SPACE = set(" \t\n\r\f\v ")
|
||
|
||
|
||
def split_source_sentences(paragraph: str):
|
||
runes = list(paragraph)
|
||
segs = []
|
||
start = i = depth = 0
|
||
n = len(runes)
|
||
while i < n:
|
||
r = runes[i]
|
||
if r in _OPEN_Q:
|
||
depth += 1
|
||
i += 1
|
||
continue
|
||
if r in _CLOSE_Q:
|
||
if depth > 0:
|
||
depth -= 1
|
||
i += 1
|
||
continue
|
||
if r not in _TERM:
|
||
i += 1
|
||
continue
|
||
j = i
|
||
cjk = False
|
||
while j < n and runes[j] in _TERM:
|
||
if runes[j] in _CJK_TERM:
|
||
cjk = True
|
||
j += 1
|
||
depth_here = depth
|
||
e = j
|
||
while e < n and runes[e] in _CLOSE_BRK:
|
||
if runes[e] in _CLOSE_Q and depth > 0:
|
||
depth -= 1
|
||
e += 1
|
||
if cjk:
|
||
boundary = depth_here == 0
|
||
else:
|
||
boundary = depth_here == 0 and (e >= n or runes[e] in _ASCII_SPACE)
|
||
# abbrev guard omitted for zh source (CJK terminators dominate; en not in S2')
|
||
if not boundary:
|
||
i = e
|
||
continue
|
||
w = e
|
||
while w < n and runes[w] in _ASCII_SPACE:
|
||
w += 1
|
||
segs.append("".join(runes[start:w]))
|
||
start = i = w
|
||
if start < n:
|
||
segs.append("".join(runes[start:]))
|
||
return segs
|
||
|
||
|
||
# --- units: split the text into atomic sentence-ish spans that TILE it -----------------------------
|
||
def units_from_text(text: str):
|
||
"""Atomic units for the DP = source sentences. Preserves whether each unit ENDS a paragraph
|
||
(\\n\\n) or a line (\\n) for boundary classification. Tiling: ''.join(u.text) == text."""
|
||
out = []
|
||
# split into paragraph blocks on blank lines, remember the separator that followed each block
|
||
# (so paragraph-end vs line-end vs sentence-end is recoverable). We tile exactly by re-attaching
|
||
# the separators to the preceding unit.
|
||
for seg, sep in _iter_blocks(text):
|
||
sents = split_source_sentences(seg)
|
||
for k, s in enumerate(sents):
|
||
ends_para = (k == len(sents) - 1) and sep == "\n\n"
|
||
ends_line = (k == len(sents) - 1) and sep == "\n"
|
||
out.append((s, ends_para, ends_line))
|
||
# re-attach the block separator to the last unit so tiling is exact
|
||
if out and sep:
|
||
s, ep, el = out[-1]
|
||
out[-1] = (s + sep, ep, el)
|
||
return out
|
||
|
||
|
||
def _iter_blocks(text: str):
|
||
"""Yield (block_text, separator) where separator is '\\n\\n' | '\\n' | ''. Tiles the input."""
|
||
i = 0
|
||
n = len(text)
|
||
buf_start = 0
|
||
while i < n:
|
||
if text[i] == "\n":
|
||
if i + 1 < n and text[i + 1] == "\n":
|
||
yield text[buf_start:i], "\n\n"
|
||
i += 2
|
||
buf_start = i
|
||
continue
|
||
else:
|
||
yield text[buf_start:i], "\n"
|
||
i += 1
|
||
buf_start = i
|
||
continue
|
||
i += 1
|
||
if buf_start < n:
|
||
yield text[buf_start:], ""
|
||
|
||
|
||
# --- boundary classification (§B1.4) ---------------------------------------------------------------
|
||
# B1 scene-shift lexicon: STRONG time/place/POV shift markers in the first chars of the NEXT unit.
|
||
# Curated generic zh markers (NOT the S2' seam text — scene-shifts must be DISCOVERED). Kept to strong
|
||
# temporal/spatial jumps + narrative shift particles; weak mid-scene connectors (顿时/忽然) excluded.
|
||
_B1_MARKERS = [
|
||
# temporal jumps
|
||
"次日", "翌日", "第二天", "第三天", "隔日", "数日后", "数天后", "几天后", "三天后", "三天三夜",
|
||
"一个星期", "一周", "几个月", "半年", "一年后", "三年后", "五年后", "多年后",
|
||
"当天", "这一天", "那一天", "这一日", "翌朝", "清晨", "早晨", "拂晓", "五更", "天亮", "天明",
|
||
"傍晚", "黄昏", "夜里", "半夜", "深夜", "入夜", "朝阳", "红日", "夕阳", "日落", "日出", "正午",
|
||
"转眼", "转眼间", "很快", "不久", "不多时", "没过多久", "光阴荏苒", "时光飞逝",
|
||
# narrative-shift particles / POV or place switch
|
||
"却说", "再说", "话说", "单说", "且说", "与此同时", "另一边", "另一方", "一方", "此时此刻",
|
||
"回到", "镜头", "视线",
|
||
]
|
||
|
||
|
||
def _b1_scene_shift(next_unit_text: str) -> bool:
|
||
head = next_unit_text.lstrip(" \n\t")[:12]
|
||
return any(m in head for m in _B1_MARKERS)
|
||
|
||
|
||
def _dialogue_boundary_penalty(prev_unit: str, next_unit: str) -> int:
|
||
"""Penalize cutting between a reply and its attribution or inside a dialogue exchange (§A.6):
|
||
if prev ends inside/with a quote and next continues speech attribution, cutting is bad."""
|
||
p = prev_unit.rstrip(" \n\t")
|
||
nx = next_unit.lstrip(" \n\t")
|
||
pen = 0
|
||
# prev ends with a closing quote AND next starts with attribution-ish (说道/道/问/答 ...) -> keep together
|
||
if p and p[-1] in _CLOSE_Q and nx[:2] in ("说道", "说", "道", "问道", "答道", "喝道", "笑道"):
|
||
pen += 1
|
||
# next starts a quote while prev did not close one cleanly (mid-exchange) — mild
|
||
if nx[:1] in _OPEN_Q and p and p[-1] in _OPEN_Q:
|
||
pen += 1
|
||
return pen
|
||
|
||
|
||
# --- TextTiling cohesion valley (B2) ---------------------------------------------------------------
|
||
def _char_bigrams(text: str):
|
||
t = [c for c in text if not c.isspace()]
|
||
return [t[k] + t[k + 1] for k in range(len(t) - 1)]
|
||
|
||
|
||
def _cosine(a: dict, b: dict) -> float:
|
||
if not a or not b:
|
||
return 0.0
|
||
dot = sum(v * b.get(k, 0) for k, v in a.items())
|
||
na = math.sqrt(sum(v * v for v in a.values()))
|
||
nb = math.sqrt(sum(v * v for v in b.values()))
|
||
return dot / (na * nb) if na and nb else 0.0
|
||
|
||
|
||
def _bag(units, lo, hi):
|
||
d = {}
|
||
for u in units[lo:hi]:
|
||
for bg in _char_bigrams(u[0]):
|
||
d[bg] = d.get(bg, 0) + 1
|
||
return d
|
||
|
||
|
||
def texttiling_depths(units, w=4):
|
||
"""Depth score of the lexical-cohesion valley at each inter-unit gap (higher = stronger boundary).
|
||
Sliding window of w units either side; cosine of char-bigram bags; depth = local valley prominence."""
|
||
n = len(units)
|
||
gaps = n - 1
|
||
sims = [0.0] * gaps
|
||
for g in range(gaps):
|
||
left = _bag(units, max(0, g + 1 - w), g + 1)
|
||
right = _bag(units, g + 1, min(n, g + 1 + w))
|
||
sims[g] = _cosine(left, right)
|
||
# depth score: for each gap, (peak_left - sim) + (peak_right - sim)
|
||
depths = [0.0] * gaps
|
||
for g in range(gaps):
|
||
lpk = sims[g]
|
||
k = g
|
||
while k > 0 and sims[k - 1] >= sims[k]:
|
||
k -= 1
|
||
lpk = max(lpk, sims[k])
|
||
rpk = sims[g]
|
||
k = g
|
||
while k < gaps - 1 and sims[k + 1] >= sims[k]:
|
||
k += 1
|
||
rpk = max(rpk, sims[k])
|
||
depths[g] = (lpk - sims[g]) + (rpk - sims[g])
|
||
return depths
|
||
|
||
|
||
# --- §B1.5 DP logical chunker ----------------------------------------------------------------------
|
||
@dataclass
|
||
class ChunkPlan:
|
||
chunks: list # list of (text, est_out, right_boundary_class)
|
||
cut_offsets: list # rune offsets in the ORIGINAL text where a cut falls (chunk starts, excl 0)
|
||
|
||
|
||
# boundary-class penalty (§B1.5): B0=0 < B1 < B2 < B3 < B4 (prefer cutting at strong boundaries)
|
||
_CLASS_PEN = {"B0": 0, "B1": 1, "B2": 3, "B3": 6, "B4": 10}
|
||
|
||
|
||
def logical_chunks(text: str, target_out: float, hard_max_out: float, min_out: float = None,
|
||
w_size=3, w_bound=60, w_dialog=200, b2_z=0.8,
|
||
corridor_lo=0.62, corridor_hi=1.12) -> ChunkPlan:
|
||
"""§B1.5 DP: minimize Σ [w_size·size_pen + w_bound·class_pen(right boundary) + w_dialog·dialogue].
|
||
Cuts only at unit (sentence) boundaries — never mid-sentence (B5 forbidden). Costs INTEGER-scaled.
|
||
target_out/hard_max_out = A0's actual est_out mean/max (size-match)."""
|
||
units = units_from_text(text)
|
||
n = len(units)
|
||
if min_out is None:
|
||
min_out = 0.45 * target_out
|
||
# per-unit est_out and cumulative rune offsets
|
||
out_of = [est_out(u[0]) for u in units]
|
||
cum_out = [0.0] * (n + 1)
|
||
cum_rune = [0] * (n + 1)
|
||
for k in range(n):
|
||
cum_out[k + 1] = cum_out[k] + out_of[k]
|
||
cum_rune[k + 1] = cum_rune[k] + len(units[k][0])
|
||
|
||
# boundary class at each inter-unit gap g (cut AFTER unit g, i.e. between unit g and g+1)
|
||
depths = texttiling_depths(units)
|
||
if depths:
|
||
mu = sum(depths) / len(depths)
|
||
sd = math.sqrt(sum((d - mu) ** 2 for d in depths) / len(depths)) or 1e-9
|
||
gap_class = []
|
||
gap_dialog = []
|
||
for g in range(n - 1):
|
||
cls = "B4" # default: sentence boundary
|
||
if units[g][2] or units[g][1]:
|
||
cls = "B3" # ends a line/paragraph
|
||
if depths and (depths[g] - mu) / sd >= b2_z:
|
||
cls = "B2" # TextTiling cohesion valley
|
||
if _b1_scene_shift(units[g + 1][0]):
|
||
cls = "B1" # scene-shift lexicon in next unit
|
||
gap_class.append(cls)
|
||
gap_dialog.append(_dialogue_boundary_penalty(units[g][0], units[g + 1][0]))
|
||
|
||
lo_band, hi_band = corridor_lo * target_out, corridor_hi * target_out
|
||
def size_pen(a, b):
|
||
"""Corridor size penalty (§B1.5 intent: size defines a BUDGET CORRIDOR; boundary class decides
|
||
WITHIN it). Flat (0) inside [lo_band, hi_band]; steep outside so chunks stay near target size."""
|
||
o = cum_out[b] - cum_out[a]
|
||
if o > hard_max_out:
|
||
return 100000 + int(40 * (o - hard_max_out)) # hard ceiling — never exceed
|
||
if o > hi_band:
|
||
return int(30 * (o - hi_band)) # over corridor
|
||
if o < lo_band:
|
||
return int(30 * (lo_band - o)) # under corridor
|
||
return 0 # inside corridor: size is free, boundary decides
|
||
|
||
INF = float("inf")
|
||
# dp[b] = min cost to segment units[0:b]; back[b] = previous cut index a
|
||
dp = [INF] * (n + 1)
|
||
back = [0] * (n + 1)
|
||
dp[0] = 0
|
||
for b in range(1, n + 1):
|
||
# try all chunk starts a < b (bounded to a window where size is plausible)
|
||
a = b - 1
|
||
while a >= 0:
|
||
o = cum_out[b] - cum_out[a]
|
||
if o > hard_max_out * 1.6 and a < b - 1:
|
||
break # window: chunks far over hardMax are never optimal
|
||
cost = dp[a] + w_size * size_pen(a, b)
|
||
if b < n: # boundary penalty on the RIGHT cut (gap b-1)
|
||
cost += w_bound * _CLASS_PEN[gap_class[b - 1]] + w_dialog * gap_dialog[b - 1]
|
||
if cost < dp[b]:
|
||
dp[b] = cost
|
||
back[b] = a
|
||
a -= 1
|
||
# reconstruct
|
||
cuts = []
|
||
b = n
|
||
while b > 0:
|
||
a = back[b]
|
||
cuts.append((a, b))
|
||
b = a
|
||
cuts.reverse()
|
||
chunks = []
|
||
cut_offsets = []
|
||
for (a, b) in cuts:
|
||
txt = "".join(units[k][0] for k in range(a, b))
|
||
rcls = gap_class[b - 1] if b < n else "B0"
|
||
chunks.append((txt, est_out(txt), rcls))
|
||
if a > 0:
|
||
cut_offsets.append(cum_rune[a])
|
||
return ChunkPlan(chunks=chunks, cut_offsets=cut_offsets)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
if "--selftest" in sys.argv:
|
||
text = Path("/home/ubuntu/books/gu-zhenren/exp15/S2prime.txt").read_text(encoding="utf-8")
|
||
units = units_from_text(text)
|
||
tiled = "".join(u[0] for u in units)
|
||
print("units:", len(units), "tiling lossless:", tiled == text)
|
||
assert tiled == text, "UNIT TILING NOT LOSSLESS"
|
||
total_out = est_out(text)
|
||
print(f"S2' total est_out={total_out:.0f} est_tokens={est_tokens(text)}")
|
||
# size sweep: how many logical chunks at a few targets
|
||
for tgt in (1500, 1797, 2200):
|
||
plan = logical_chunks(text, target_out=tgt, hard_max_out=tgt * 1.25)
|
||
outs = [c[1] for c in plan.chunks]
|
||
classes = [c[2] for c in plan.chunks[:-1]] # right-boundary classes (excl last)
|
||
from collections import Counter
|
||
print(f" target_out={tgt}: {len(plan.chunks)} chunks "
|
||
f"mean_out={sum(outs)/len(outs):.0f} min={min(outs):.0f} max={max(outs):.0f} "
|
||
f"cut-classes={dict(Counter(classes))}")
|
||
print("chunker selftest OK")
|