348ce10b1d
chunk_text/hard_split/split_paragraphs + 2400/600-char constants move out of documents_ingest.chunk_embed into kb_mail.chunking so the upcoming jobs/mail-body-ingest shares the exact same chunker instead of a copy-pasted drift (the phase-4 lesson for retrieval.py -> packages/kb-retrieval, applied again). chunk_embed.py re-exports unchanged; zero behavior change, tests moved 1:1 alongside the code (docs/kb/modules/05-faza-mailowa-plan.md, §3).
86 lines
3 KiB
Python
86 lines
3 KiB
Python
"""Shared chunker — module 5, phase mailowa, plan step 0 (docs/kb/modules/05-faza-mailowa-plan.md,
|
|
§3). Moved 1:1 out of `jobs/documents_ingest/chunk_embed.py` so both the paperless job and
|
|
`jobs/mail-body-ingest` chunk against the exact same tested logic instead of two implementations
|
|
drifting apart (the phase-4 lesson for `retrieval.py` -> `packages/kb-retrieval`, applied again).
|
|
|
|
No new dependencies here — `is_ocr_junk` (OCR-corpus calibrated) and the embed client stay in
|
|
their own modules; this file is chunking only.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
|
|
# ~600 tok/chunk, ~150 tok overlap. No local bge-m3 tokenizer available offline, so tokens are
|
|
# approximated at ~4 chars/token (common heuristic) -> ~2400 chars target, ~600 chars overlap.
|
|
# Paragraph boundaries are preferred; a hard character-based sliding window is only used when a
|
|
# single paragraph alone exceeds the target.
|
|
CHARS_PER_TOKEN = 4
|
|
TARGET_TOKENS = 600
|
|
OVERLAP_TOKENS = 150
|
|
TARGET_CHARS = TARGET_TOKENS * CHARS_PER_TOKEN
|
|
OVERLAP_CHARS = OVERLAP_TOKENS * CHARS_PER_TOKEN
|
|
|
|
_PARA_SPLIT = re.compile(r"\n\s*\n")
|
|
|
|
|
|
def split_paragraphs(text: str) -> list[str]:
|
|
"""Split on blank-line boundaries; drop empty fragments."""
|
|
return [p.strip() for p in _PARA_SPLIT.split(text) if p.strip()]
|
|
|
|
|
|
def hard_split(text: str, size: int = TARGET_CHARS, overlap: int = OVERLAP_CHARS) -> list[str]:
|
|
"""Fixed-size sliding window over raw characters — fallback for a paragraph too big to fit."""
|
|
if overlap >= size:
|
|
# step = size - overlap would be <= 0 below, so `start` would never advance past `n`
|
|
# and the while loop would spin forever, growing `chunks` without bound.
|
|
raise ValueError(f"chunk_overlap ({overlap}) must be smaller than chunk_size ({size})")
|
|
if len(text) <= size:
|
|
return [text]
|
|
step = size - overlap
|
|
chunks = []
|
|
start = 0
|
|
n = len(text)
|
|
while start < n:
|
|
end = min(start + size, n)
|
|
chunks.append(text[start:end])
|
|
if end == n:
|
|
break
|
|
start += step
|
|
return chunks
|
|
|
|
|
|
def chunk_text(text: str, size: int = TARGET_CHARS, overlap: int = OVERLAP_CHARS) -> list[str]:
|
|
"""Paragraph-preferring chunker with a hard character-fallback.
|
|
|
|
Empty/whitespace-only text returns []. A document shorter than one chunk returns exactly
|
|
one chunk (the full text).
|
|
"""
|
|
text = (text or "").strip()
|
|
if not text:
|
|
return []
|
|
|
|
paragraphs = split_paragraphs(text)
|
|
|
|
chunks: list[str] = []
|
|
buffer = ""
|
|
for para in paragraphs:
|
|
if len(para) > size:
|
|
if buffer:
|
|
chunks.append(buffer)
|
|
buffer = ""
|
|
chunks.extend(hard_split(para, size, overlap))
|
|
continue
|
|
|
|
candidate = f"{buffer}\n\n{para}" if buffer else para
|
|
if len(candidate) <= size:
|
|
buffer = candidate
|
|
continue
|
|
|
|
chunks.append(buffer)
|
|
tail = buffer[-overlap:] if buffer else ""
|
|
buffer = f"{tail}\n\n{para}" if tail else para
|
|
|
|
if buffer:
|
|
chunks.append(buffer)
|
|
return chunks
|