File size: 4,355 Bytes
30e6218 | 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 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | """Offset-preserving character chunker.
Limits are in *characters* of the normalized extracted text. Every chunk
records its ``[char_start, char_end)`` span, the page range it covers (PDF
only), and the nearest preceding Markdown heading (Markdown only). The
algorithm always makes progress, so it terminates for any input, and it never
emits whitespace-only chunks.
"""
from __future__ import annotations
import hashlib
from dataclasses import dataclass
from typing import List, Optional, Sequence
from ..utils.document_loader import PageSpan, markdown_headings
from .types import IngestionError, SourceLocation
_BOUNDARIES: Sequence[str] = ("\n\n", "\n", ". ", "? ", "! ", "; ", ", ", " ")
@dataclass(frozen=True)
class Chunk:
ordinal: int
text: str
location: SourceLocation
@property
def text_hash(self) -> str:
return hashlib.sha256(self.text.encode("utf-8")).hexdigest()[:16]
def _find_boundary(text: str, lo: int, hi: int) -> Optional[int]:
"""Best split point in ``(lo, hi]`` preferring stronger separators."""
for sep in _BOUNDARIES:
idx = text.rfind(sep, lo, hi)
if idx != -1 and idx + len(sep) > lo:
return idx + len(sep)
return None
def _page_for(offset: int, pages: Sequence[PageSpan]) -> Optional[int]:
for span in pages:
if span.char_start <= offset < span.char_end:
return span.page_number
# Offsets that land in a separator between pages belong to the next page.
for span in pages:
if offset < span.char_start:
return span.page_number
return pages[-1].page_number if pages else None
def _section_for(offset: int, headings: Sequence[dict]) -> Optional[str]:
current = None
for heading in headings:
if int(heading["start"]) <= offset:
current = str(heading["title"])
else:
break
return current
def chunk_text(
text: str,
*,
chunk_size: int,
chunk_overlap: int,
max_chunks: int,
pages: Sequence[PageSpan] = (),
extension: str = ".txt",
) -> List[Chunk]:
"""Split ``text`` into overlapping chunks with provenance."""
if chunk_size <= 0:
raise ValueError("chunk_size must be positive")
if chunk_overlap < 0 or chunk_overlap >= chunk_size:
raise ValueError("chunk_overlap must be >= 0 and < chunk_size")
if max_chunks <= 0:
raise ValueError("max_chunks must be positive")
headings = markdown_headings(text) if extension == ".md" else []
length = len(text)
chunks: List[Chunk] = []
start = 0
min_step = max(1, chunk_size - chunk_overlap)
while start < length:
# Skip leading whitespace without losing offset accuracy.
while start < length and text[start].isspace():
start += 1
if start >= length:
break
end = min(start + chunk_size, length)
if end < length:
boundary = _find_boundary(text, start + chunk_size // 2, end)
if boundary is not None and boundary > start:
end = boundary
piece = text[start:end]
stripped = piece.rstrip()
real_end = start + len(stripped)
if stripped.strip():
if len(chunks) >= max_chunks:
raise IngestionError(
f"Document produces more than {max_chunks} chunks",
code="too_many_chunks",
)
location = SourceLocation(
char_start=start,
char_end=real_end,
page=_page_for(start, pages) if pages else None,
page_end=_page_for(max(real_end - 1, start), pages) if pages else None,
section=_section_for(start, headings) if headings else None,
)
chunks.append(Chunk(ordinal=len(chunks), text=stripped, location=location))
if end >= length:
break
next_start = end - chunk_overlap
# Move the overlap start forward to a whitespace boundary if one exists
# inside the overlap region, so chunks do not start mid-word.
probe = text.find(" ", next_start, end)
if probe != -1 and probe + 1 < end:
next_start = probe + 1
if next_start <= start:
next_start = start + min_step
start = next_start
return chunks
|