File size: 2,784 Bytes
325b94c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | # chunking.py
import re
from typing import List, Tuple
def approx_token_count(text: str) -> int:
"""
Rough token estimate for chunk sizing (heuristic).
Keep this simple but consistent: ~4 chars per token.
"""
return max(1, len(text) // 4)
def chunk_pages(
pages: List[str],
target_tokens: int = 520,
overlap_tokens: int = 80,
) -> List[Tuple[str, int, int]]:
"""
Split pages (list[str]) into overlapping chunks.
Returns list of tuples: (chunk_text, page_start, page_end)
Overlap is implemented at paragraph level (keeps page ranges correct).
"""
chunks: List[Tuple[str, int, int]] = []
buffer: List[str] = []
buffer_pages: List[int] = []
buffer_tokens = 0
def make_chunk():
nonlocal buffer, buffer_pages, buffer_tokens
if not buffer:
return
chunk_text = "\n\n".join(buffer).strip()
page_start = min(buffer_pages)
page_end = max(buffer_pages)
chunks.append((chunk_text, page_start, page_end))
for page_idx, page in enumerate(pages, start=1):
paragraphs = [p.strip() for p in re.split(r"\n\s*\n", page) if p.strip()]
for para in paragraphs:
t = approx_token_count(para)
# If single paragraph exceeds target, create it as its own chunk
if buffer_tokens == 0 and t > target_tokens:
# make chunk with this large paragraph alone
buffer = [para]
buffer_pages = [page_idx]
buffer_tokens = t
make_chunk()
buffer, buffer_pages, buffer_tokens = [], [], 0
continue
# If adding this paragraph would exceed target, flush current chunk
if buffer_tokens + t > target_tokens:
make_chunk()
# prepare overlap: keep tail paragraphs whose tokens sum >= overlap_tokens
tail_buffer: List[str] = []
tail_pages: List[int] = []
tail_tokens = 0
# iterate buffer in reverse to pick tail paragraphs
for p, p_pg in zip(reversed(buffer), reversed(buffer_pages)):
pt = approx_token_count(p)
tail_buffer.insert(0, p)
tail_pages.insert(0, p_pg)
tail_tokens += pt
if tail_tokens >= overlap_tokens:
break
buffer = tail_buffer
buffer_pages = tail_pages
buffer_tokens = tail_tokens
# append current paragraph
buffer.append(para)
buffer_pages.append(page_idx)
buffer_tokens += t
# final flush
if buffer:
make_chunk()
return chunks
|