118 lines
4.7 KiB
Python
118 lines
4.7 KiB
Python
"""Comprehensive round-trip test for the forge (refactored) chunker.
|
|
|
|
Covers:
|
|
1. Repetitive sentence-grouped text (the probe case)
|
|
2. Long single sentence -> triggers _hard_split_long
|
|
3. Mixed: some sentences fit, some need hard split
|
|
4. Edge cases: empty, whitespace-only, single chunk, no overlap
|
|
5. Many random-ish inputs with varying params
|
|
"""
|
|
import sys
|
|
import random
|
|
import string
|
|
sys.path.insert(0, "/root/.hermes/wizard-outputs/forge")
|
|
import chunker
|
|
|
|
PASS = 0
|
|
FAIL = 0
|
|
|
|
def check(name, text, max_tokens=256, overlap_tokens=32):
|
|
global PASS, FAIL
|
|
if not text.strip():
|
|
expected = ""
|
|
else:
|
|
expected = text.strip()
|
|
cs = chunker.chunk_text(text, max_tokens=max_tokens, overlap_tokens=overlap_tokens)
|
|
r = chunker.reconstruct(cs, overlap_tokens)
|
|
ok = r == expected
|
|
status = "PASS" if ok else "FAIL"
|
|
if ok:
|
|
PASS += 1
|
|
else:
|
|
FAIL += 1
|
|
print(f"[{status}] {name}")
|
|
print(f" orig_len={len(expected)} recon_len={len(r)}")
|
|
for i in range(min(len(r), len(expected))):
|
|
if r[i] != expected[i]:
|
|
print(f" FIRST DIFF at char {i}")
|
|
print(f" orig: {expected[max(0,i-30):i+30]!r}")
|
|
print(f" recon: {r[max(0,i-30):i+30]!r}")
|
|
break
|
|
else:
|
|
print(f" common prefix OK but lengths differ: {len(r)} vs {len(expected)}")
|
|
print(f" orig tail: {expected[-30:]!r}")
|
|
print(f" recon tail: {r[-30:]!r}")
|
|
print(f" nchunks={len(cs)}")
|
|
# show if any chunk exceeded max_tokens
|
|
for i, c in enumerate(cs):
|
|
# strip the overlap tail for token check
|
|
actual = c
|
|
if chunker._OVERLAP_SENTINEL in actual:
|
|
actual = actual.partition(chunker._OVERLAP_SENTINEL)[2]
|
|
est = chunker.estimate_tokens(actual)
|
|
if est > max_tokens:
|
|
print(f" chunk[{i}] exceeds max_tokens: est={est} > {max_tokens}, content={actual[:60]!r}")
|
|
return ok
|
|
|
|
|
|
# Case 1: repetitive sentence-grouped text
|
|
base = "The quick brown fox jumps over the lazy dog. Pack my box with five dozen liquor jugs. How vexingly quick daft zebras jump. "
|
|
text = (base + "Variant sentence here. ") * 20
|
|
check("repetitive-sentences", text, max_tokens=64, overlap_tokens=8)
|
|
|
|
# Case 2: long single sentence -> hard split
|
|
long_sent = "Word" + " word" * 200 + " end." # ~1000 chars, way over 80
|
|
check("long-sentence-hard-split", long_sent, max_tokens=20, overlap_tokens=4)
|
|
|
|
# Case 2b: long single sentence with sentence-ending punctuation
|
|
long_sent2 = " ".join(["Sentence"] + ["word"]*300 + ["."])
|
|
check("long-sentence-hard-split-v2", long_sent2, max_tokens=30, overlap_tokens=6)
|
|
|
|
# Case 3: mixed - some short, some long sentences
|
|
mixed = "Short one. " + long_sent + " Another short. " + long_sent2 + " Final."
|
|
check("mixed-sizes", mixed, max_tokens=20, overlap_tokens=4)
|
|
|
|
# Case 4: edge cases
|
|
check("empty-string", "", max_tokens=64, overlap_tokens=8)
|
|
check("whitespace-only", " \n\t ", max_tokens=64, overlap_tokens=8)
|
|
check("single-sentence", "Just one sentence here.", max_tokens=64, overlap_tokens=8)
|
|
check("two-sentences-no-overlap", "First sentence. Second sentence.", max_tokens=256, overlap_tokens=0)
|
|
check("small-overlap", "First sentence. Second sentence.", max_tokens=64, overlap_tokens=2)
|
|
|
|
# Case 5: multiple sentences exactly at boundary
|
|
exactly = " ".join(f"Sentence number {i}." for i in range(10))
|
|
check("boundary-fit", exactly, max_tokens=10, overlap_tokens=2)
|
|
|
|
# Case 6: text with special chars, multiple spaces
|
|
special = "Hello world. How are you? I am fine. " * 5
|
|
check("multi-space", special, max_tokens=32, overlap_tokens=4)
|
|
|
|
# Case 7: text starting/ending with whitespace
|
|
ws_text = " The quick brown fox. The lazy dog. "
|
|
check("leading-trailing-ws", ws_text, max_tokens=64, overlap_tokens=8)
|
|
|
|
# Case 8: single long sentence that needs hard split, with high overlap
|
|
long3 = "A" + " b" * 500 + " end."
|
|
check("long-with-high-overlap", long3, max_tokens=16, overlap_tokens=8)
|
|
|
|
# Case 9: random text fuzz
|
|
random.seed(42)
|
|
words = ["alpha", "beta", "gamma", "delta", "epsilon", "zeta", "eta", "theta"]
|
|
for trial in range(20):
|
|
sents = []
|
|
for _ in range(random.randint(3, 15)):
|
|
nwords = random.randint(3, 50)
|
|
s = " ".join(random.choice(words) for _ in range(nwords))
|
|
sents.append(s + ".")
|
|
text = " ".join(sents)
|
|
check(f"random-trial-{trial}", text, max_tokens=random.randint(16, 64),
|
|
overlap_tokens=random.randint(0, 8))
|
|
|
|
# Case 10: text with newlines
|
|
nl_text = "First paragraph.\n\nSecond paragraph here.\nThird line."
|
|
check("newlines", nl_text, max_tokens=64, overlap_tokens=8)
|
|
|
|
print(f"\n=== Results: {PASS} passed, {FAIL} failed ===")
|
|
if FAIL:
|
|
sys.exit(1)
|