362 lines
14 KiB
Go
362 lines
14 KiB
Go
package pipeline
|
||
|
||
import (
|
||
"reflect"
|
||
"strings"
|
||
"testing"
|
||
"unicode"
|
||
"unicode/utf8"
|
||
)
|
||
|
||
// testSeg is the ratified WS2 output-token budget (zh-ru defaults) used by the chunker tests.
|
||
func testSeg() SegBudget {
|
||
return SegBudget{DraftBudgetOut: 1797, EditCeilingOut: 3200, FertCJK: 1.1978, FertOther: 0.3852}
|
||
}
|
||
|
||
// stripMeta zeroes the WS2 additive fields (EstOut/OversizedSentence/EditUnitID) so a
|
||
// {Chapter,ChunkIdx,Text}-level DeepEqual still expresses the intent of the pre-WS2 assertions.
|
||
func stripMeta(cs []Chunk) []Chunk {
|
||
out := make([]Chunk, len(cs))
|
||
for i, c := range cs {
|
||
out[i] = Chunk{Chapter: c.Chapter, ChunkIdx: c.ChunkIdx, Text: c.Text}
|
||
}
|
||
return out
|
||
}
|
||
|
||
func TestSplitChunksSingleParagraph(t *testing.T) {
|
||
got := SplitChunks([]string{"静かな図書館の朝。"}, testSeg())
|
||
want := []Chunk{{Chapter: 1, ChunkIdx: 0, Text: "静かな図書館の朝。"}}
|
||
if !reflect.DeepEqual(stripMeta(got), want) {
|
||
t.Fatalf("single paragraph = %+v, want %+v", got, want)
|
||
}
|
||
// WS2: the additive fields are populated — est_out>0 and a single unit id.
|
||
if got[0].EstOut <= 0 || got[0].EditUnitID != 0 || got[0].OversizedSentence {
|
||
t.Fatalf("WS2 fields = est_out %.3f unit %d oversized %v", got[0].EstOut, got[0].EditUnitID, got[0].OversizedSentence)
|
||
}
|
||
}
|
||
|
||
func TestSplitChunksChapters(t *testing.T) {
|
||
got := SplitChunks([]string{"Глава один.", "Глава два.", "Глава три."}, testSeg())
|
||
want := []Chunk{
|
||
{Chapter: 1, ChunkIdx: 0, Text: "Глава один."},
|
||
{Chapter: 2, ChunkIdx: 0, Text: "Глава два."},
|
||
{Chapter: 3, ChunkIdx: 0, Text: "Глава три."},
|
||
}
|
||
if !reflect.DeepEqual(stripMeta(got), want) {
|
||
t.Fatalf("chapters = %+v, want %+v", got, want)
|
||
}
|
||
// WS2: each chapter is its own edit unit → distinct monotone EditUnitIDs 0,1,2.
|
||
for i, c := range got {
|
||
if c.EditUnitID != i {
|
||
t.Fatalf("chapter %d chunk edit_unit_id = %d, want %d (each short chapter is its own unit)", c.Chapter, c.EditUnitID, i)
|
||
}
|
||
}
|
||
}
|
||
|
||
func TestSplitChunksPacksToTarget(t *testing.T) {
|
||
// Two paragraphs that together exceed the OUTPUT-token budget (each ~1156 out-tokens for
|
||
// cyrillic: 3000 chars × 0.3852) must split into two draft chunks, each kept whole.
|
||
p1 := strings.Repeat("а", 3000)
|
||
p2 := strings.Repeat("б", 3000)
|
||
got := SplitChunks([]string{p1 + "\n\n" + p2}, testSeg())
|
||
want := []Chunk{
|
||
{Chapter: 1, ChunkIdx: 0, Text: p1},
|
||
{Chapter: 1, ChunkIdx: 1, Text: p2},
|
||
}
|
||
if !reflect.DeepEqual(stripMeta(got), want) {
|
||
t.Fatalf("packing = %d chunks, want 2 whole paragraphs; got %+v", len(got), summarize(got))
|
||
}
|
||
// WS2: the two ~1156-out chunks (2311 together ≤ 3200) group into ONE edit unit.
|
||
if got[0].EditUnitID != got[1].EditUnitID {
|
||
t.Fatalf("two sub-ceiling draft chunks must share an edit unit, got ids %d and %d", got[0].EditUnitID, got[1].EditUnitID)
|
||
}
|
||
}
|
||
|
||
func TestSplitChunksPacksSmallParagraphsTogether(t *testing.T) {
|
||
got := SplitChunks([]string{"Абзац один.\n\nАбзац два.\n\nАбзац три."}, testSeg())
|
||
if len(got) != 1 {
|
||
t.Fatalf("small paragraphs must pack into one chunk, got %d: %+v", len(got), summarize(got))
|
||
}
|
||
if got[0].Text != "Абзац один.\n\nАбзац два.\n\nАбзац три." {
|
||
t.Fatalf("packed text = %q", got[0].Text)
|
||
}
|
||
}
|
||
|
||
func TestSplitChunksOversizeParagraphSplitsAtSentences(t *testing.T) {
|
||
// One paragraph far over the output budget, built of many whole sentences, must
|
||
// split into >1 chunk AND never cut a sentence: every original sentence must
|
||
// survive intact and in order across the chunks.
|
||
sentence := strings.Repeat("слово ", 40) + "конец." // ~79 out-tokens each
|
||
var sb strings.Builder
|
||
for i := 0; i < 40; i++ { // ~3160 out-tokens total ≫ 1797
|
||
sb.WriteString(sentence)
|
||
sb.WriteByte(' ')
|
||
}
|
||
para := strings.TrimSpace(sb.String())
|
||
got := SplitChunks([]string{para}, testSeg())
|
||
if len(got) < 2 {
|
||
t.Fatalf("oversize paragraph must split into >1 chunk, got %d", len(got))
|
||
}
|
||
wantSentences := trimAll(splitSourceSentences(para))
|
||
var gotSentences []string
|
||
for _, c := range got {
|
||
gotSentences = append(gotSentences, trimAll(splitSourceSentences(c.Text))...)
|
||
}
|
||
if !reflect.DeepEqual(gotSentences, wantSentences) {
|
||
t.Fatalf("sentence integrity broken:\n got %d sentences\nwant %d sentences", len(gotSentences), len(wantSentences))
|
||
}
|
||
}
|
||
|
||
func TestSplitChunksShortLinesPackByTrueTokenBudget(t *testing.T) {
|
||
// Many short dialogue paragraphs must pack by TRUE output-token content (est_out on the
|
||
// joined text), not by summing a per-unit floor. 150 lines of 「はい。」 are ~533 out-tokens,
|
||
// so they belong in ONE chunk (a per-unit floor would wrongly flush into sub-500-char pieces).
|
||
var paras []string
|
||
for i := 0; i < 150; i++ {
|
||
paras = append(paras, "「はい。」")
|
||
}
|
||
got := SplitChunks([]string{strings.Join(paras, "\n\n")}, testSeg())
|
||
if len(got) != 1 {
|
||
t.Fatalf("short lines must pack by true token budget into 1 chunk, got %d chunks", len(got))
|
||
}
|
||
if got[0].EstOut > testSeg().DraftBudgetOut {
|
||
t.Fatalf("packed chunk overshoots the budget: est_out %.1f", got[0].EstOut)
|
||
}
|
||
}
|
||
|
||
func TestSplitChunksOversizeSentenceIsOwnChunk(t *testing.T) {
|
||
// A single sentence (no internal terminator) larger than the budget can only be its own
|
||
// chunk — never split — and is flagged OversizedSentence (passthrough, WS2 §2б).
|
||
big := strings.Repeat("ы", 6000) // ~2311 out-tokens, one un-terminated run
|
||
got := SplitChunks([]string{big}, testSeg())
|
||
if len(got) != 1 || got[0].Text != big {
|
||
t.Fatalf("an oversize sentence must be one un-split chunk, got %+v", summarize(got))
|
||
}
|
||
if !got[0].OversizedSentence {
|
||
t.Fatalf("a single over-budget sentence must be flagged OversizedSentence")
|
||
}
|
||
}
|
||
|
||
func TestSplitChunksDropsEmpties(t *testing.T) {
|
||
// A blank/whitespace-only chapter vanishes and does NOT consume a chapter number.
|
||
got := SplitChunks([]string{"A", " \n\n ", "B"}, testSeg())
|
||
want := []Chunk{
|
||
{Chapter: 1, ChunkIdx: 0, Text: "A"},
|
||
{Chapter: 2, ChunkIdx: 0, Text: "B"},
|
||
}
|
||
if !reflect.DeepEqual(stripMeta(got), want) {
|
||
t.Fatalf("empties = %+v, want %+v", got, want)
|
||
}
|
||
}
|
||
|
||
func TestSplitChunksDeterministic(t *testing.T) {
|
||
chapters := []string{
|
||
"Абзац.\n\nЕщё абзац.",
|
||
strings.Repeat("слово ", 400) + "конец.",
|
||
}
|
||
a, b := SplitChunks(chapters, testSeg()), SplitChunks(chapters, testSeg())
|
||
if !reflect.DeepEqual(a, b) {
|
||
t.Fatal("SplitChunks must be deterministic")
|
||
}
|
||
}
|
||
|
||
func TestSplitChunksNeverSplitsAbbreviation(t *testing.T) {
|
||
// Two long en paragraphs each ending mid-flow with "Mr. Smith" must not chunk
|
||
// between "Mr." and "Smith" — the abbreviation guard keeps the sentence whole.
|
||
p := strings.Repeat("The road went on and on. ", 100) + "It was Mr. Smith who arrived."
|
||
got := SplitChunks([]string{p}, testSeg())
|
||
for _, c := range got {
|
||
if strings.HasSuffix(strings.TrimSpace(c.Text), "Mr.") {
|
||
t.Fatalf("chunk ended at an abbreviation 'Mr.' — a sentence was cut: %q", c.Text)
|
||
}
|
||
}
|
||
}
|
||
|
||
// TestSplitChunksEditUnitGroupsWholeChunks pins the WS2 decoupling: a chapter whose draft chunks
|
||
// together exceed EditCeilingOut is grouped into >1 edit unit at DRAFT-CHUNK boundaries (never
|
||
// splitting a chunk), so every chunk of a unit reconstructs the unit's draft losslessly.
|
||
func TestSplitChunksEditUnitGroupsWholeChunks(t *testing.T) {
|
||
// Three ~1156-out paragraphs (0.3852×3000): draft chunks {p1},{p2},{p3} (each ~1156, two
|
||
// together 2311 ≤ 1797? no — 2311>1797 so each is its own draft chunk). Edit ceiling 3200:
|
||
// group whole chunks → {p1,p2}=2311 ≤3200, +p3=3467>3200 → unit0={p1,p2}, unit1={p3}.
|
||
p := strings.Repeat("а", 3000)
|
||
q := strings.Repeat("б", 3000)
|
||
r := strings.Repeat("в", 3000)
|
||
got := SplitChunks([]string{p + "\n\n" + q + "\n\n" + r}, testSeg())
|
||
if len(got) != 3 {
|
||
t.Fatalf("want 3 draft chunks, got %d", len(got))
|
||
}
|
||
ids := []int{got[0].EditUnitID, got[1].EditUnitID, got[2].EditUnitID}
|
||
if ids[0] != ids[1] || ids[1] == ids[2] {
|
||
t.Fatalf("edit-unit grouping = %v, want first two together and the third separate", ids)
|
||
}
|
||
// Every draft chunk belongs to exactly one unit; a unit's chunks are contiguous.
|
||
for i := 1; i < len(got); i++ {
|
||
if got[i].EditUnitID < got[i-1].EditUnitID {
|
||
t.Fatalf("edit unit ids must be monotone, got %v", ids)
|
||
}
|
||
}
|
||
}
|
||
|
||
// --- source sentence splitter (zh/ja/en) ---------------------------------------
|
||
|
||
func TestSplitSourceSentencesTilesLosslessly(t *testing.T) {
|
||
// The segments must TILE the input: concatenation reproduces it byte-for-byte.
|
||
inputs := []string{
|
||
"静かな朝。彼は歩いた。「止まれ！」と叫んだ。",
|
||
"他说。“不要！”然后离开了。",
|
||
"Hello world. This is a test! Really? Yes.",
|
||
"Mr. Smith met Dr. J. R. R. Brown. They talked.",
|
||
"末尾に句点なし",
|
||
"",
|
||
}
|
||
for _, in := range inputs {
|
||
segs := splitSourceSentences(in)
|
||
if got := strings.Join(segs, ""); got != in {
|
||
t.Fatalf("tiling broke for %q: rejoined to %q", in, got)
|
||
}
|
||
}
|
||
}
|
||
|
||
func TestSplitSourceSentencesJapanese(t *testing.T) {
|
||
// CJK terminators split zero-width; a closing quote is absorbed into its sentence.
|
||
got := trimAll(splitSourceSentences("静かな朝。「止まれ！」と彼は言った。"))
|
||
want := []string{"静かな朝。", "「止まれ！」と彼は言った。"}
|
||
if !reflect.DeepEqual(got, want) {
|
||
t.Fatalf("ja split = %q, want %q", got, want)
|
||
}
|
||
}
|
||
|
||
func TestSplitSourceSentencesEnglish(t *testing.T) {
|
||
got := trimAll(splitSourceSentences(`He said "Stop!" She ran. Done.`))
|
||
want := []string{`He said "Stop!"`, "She ran.", "Done."}
|
||
if !reflect.DeepEqual(got, want) {
|
||
t.Fatalf("en split = %q, want %q", got, want)
|
||
}
|
||
}
|
||
|
||
func TestSplitSourceSentencesEllipsisNotOverSplitInCJK(t *testing.T) {
|
||
// "……" mid-CJK-sentence must not split; the trailing 。 still ends it.
|
||
got := trimAll(splitSourceSentences("そう……だった。次へ。"))
|
||
want := []string{"そう……だった。", "次へ。"}
|
||
if !reflect.DeepEqual(got, want) {
|
||
t.Fatalf("ellipsis split = %q, want %q", got, want)
|
||
}
|
||
}
|
||
|
||
func TestSplitSourceSentencesASCIIRespectsQuoteDepth(t *testing.T) {
|
||
// ASCII terminators inside a tracked quote must NOT split (like the CJK branch):
|
||
// 「Stop. Now.」 is one dialogue unit, not two sentences (external-review #4).
|
||
got := trimAll(splitSourceSentences("「Stop. Now.」次へ。"))
|
||
want := []string{"「Stop. Now.」次へ。"}
|
||
if !reflect.DeepEqual(got, want) {
|
||
t.Fatalf("ASCII-in-quote split = %q, want %q (must not over-split dialogue)", got, want)
|
||
}
|
||
// Sanity: outside quotes ASCII still splits.
|
||
if n := len(trimAll(splitSourceSentences("Stop. Now."))); n != 2 {
|
||
t.Fatalf("plain ASCII must still split into 2, got %d", n)
|
||
}
|
||
}
|
||
|
||
func TestSplitSourceSentencesAbbreviations(t *testing.T) {
|
||
// Abbreviations and single-letter initials do not end a sentence.
|
||
got := trimAll(splitSourceSentences("Dr. J. R. R. Tolkien wrote it. The end."))
|
||
want := []string{"Dr. J. R. R. Tolkien wrote it.", "The end."}
|
||
if !reflect.DeepEqual(got, want) {
|
||
t.Fatalf("abbrev split = %q, want %q", got, want)
|
||
}
|
||
}
|
||
|
||
// --- property/fuzz (external-review #7: would have caught #3) ----------------
|
||
|
||
// FuzzSplitSourceSentencesTiles asserts the segmenter is lossless: its segments
|
||
// always concatenate back to the input's rune sequence, and byte-exactly when the
|
||
// input is valid UTF-8 (the production invariant after NormalizeSource). A path that
|
||
// dropped/reordered/replaced a rune — like the size-dependent []rune→U+FFFD bug #3 —
|
||
// fails this over random input.
|
||
func FuzzSplitSourceSentencesTiles(f *testing.F) {
|
||
for _, s := range []string{
|
||
"静かな朝。「止まれ！」と言った。", "Hello. World!", "Mr. Smith went. Home.",
|
||
"そう……だった。", "「Stop. Now.」", "abc\xe9def", "",
|
||
} {
|
||
f.Add(s)
|
||
}
|
||
f.Fuzz(func(t *testing.T, s string) {
|
||
segs := splitSourceSentences(s)
|
||
joined := strings.Join(segs, "")
|
||
if joined != string([]rune(s)) {
|
||
t.Fatalf("tiling lost/reordered runes: input=%q rejoined=%q", s, joined)
|
||
}
|
||
if utf8.ValidString(s) && joined != s {
|
||
t.Fatalf("valid-UTF-8 tiling not byte-exact: %q → %q", s, joined)
|
||
}
|
||
})
|
||
}
|
||
|
||
// FuzzNormalizeSourceValidUTF8 guards fix #3 directly: NormalizeSource always yields
|
||
// valid UTF-8, so the chunker never sees a stray byte the two paths would treat
|
||
// differently.
|
||
func FuzzNormalizeSourceValidUTF8(f *testing.F) {
|
||
for _, s := range []string{"abc\xe9def", "正常なテキスト", "\xff\xfe", ""} {
|
||
f.Add(s)
|
||
}
|
||
f.Fuzz(func(t *testing.T, s string) {
|
||
if out := NormalizeSource(s); !utf8.ValidString(out) {
|
||
t.Fatalf("NormalizeSource must return valid UTF-8, invalid for %q", s)
|
||
}
|
||
})
|
||
}
|
||
|
||
// FuzzSplitChunksPreservesContent asserts chunking loses no non-space content: the
|
||
// count of non-space runes across all chunks equals that of the normalized source.
|
||
func FuzzSplitChunksPreservesContent(f *testing.F) {
|
||
for _, s := range []string{
|
||
"静かな朝。\n\n「止まれ！」と彼は言った。次へ。", strings.Repeat("слово ", 500), "abc\xe9def", "",
|
||
} {
|
||
f.Add(s)
|
||
}
|
||
f.Fuzz(func(t *testing.T, s string) {
|
||
norm := NormalizeSource(s)
|
||
want := countNonSpace(norm)
|
||
got := 0
|
||
for _, c := range SplitChunks([]string{norm}, testSeg()) {
|
||
got += countNonSpace(c.Text)
|
||
}
|
||
if got != want {
|
||
t.Fatalf("chunking lost non-space content: got %d want %d (input %q)", got, want, s)
|
||
}
|
||
})
|
||
}
|
||
|
||
func countNonSpace(s string) int {
|
||
n := 0
|
||
for _, r := range s {
|
||
if !unicode.IsSpace(r) {
|
||
n++
|
||
}
|
||
}
|
||
return n
|
||
}
|
||
|
||
// trimAll TrimSpaces each segment and drops empties (test convenience).
|
||
func trimAll(segs []string) []string {
|
||
var out []string
|
||
for _, s := range segs {
|
||
if t := strings.TrimSpace(s); t != "" {
|
||
out = append(out, t)
|
||
}
|
||
}
|
||
return out
|
||
}
|
||
|
||
func summarize(cs []Chunk) []string {
|
||
out := make([]string, len(cs))
|
||
for i, c := range cs {
|
||
txt := c.Text
|
||
if len(txt) > 12 {
|
||
txt = txt[:12] + "…"
|
||
}
|
||
out[i] = strings.TrimSpace(txt)
|
||
}
|
||
return out
|
||
}
|