d1348cac3f
Shorten and condense comments across the RAG backend, frontend, and tests for readability. Comment text only; no code, strings, identifiers, or logic changed. License headers and lint/type pragmas are preserved.
307 lines
9.8 KiB
Python
307 lines
9.8 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
from dataclasses import dataclass
|
|
from typing import Callable
|
|
|
|
from .parsers import ParsedPage
|
|
|
|
# Match "Figure 1:", "Figure 1.2:", "Fig. 3.", "Table 4:" etc. at line-start,
|
|
# tolerating leading bold markers. Breaks chunks BEFORE such captions so the
|
|
# caption starts its own chunk — dense embeddings pool over the whole chunk, so
|
|
# figure refs buried at the end get diluted by surrounding body text.
|
|
_FIGURE_BOUNDARY_RE = re.compile(
|
|
# Number forms: "1", "12", "1.2", "B.1" (appendix); bold wrappers tolerated.
|
|
r"^\**(?:Figure|Fig\.|Table|Tab\.)\s+[A-Z]?\.?\d+(?:\.\d+)?\**[\.:]",
|
|
re.MULTILINE | re.IGNORECASE,
|
|
)
|
|
|
|
|
|
def _split_at_figure_boundaries(text: str) -> list[str]:
|
|
"""Split markdown at the start of each figure / table caption.
|
|
|
|
Each segment starts with either the original text head or a "Figure N:" /
|
|
"Table N:" line, so the caption anchors the embedding of the chunk it
|
|
lands in. Returns the original text as a single-element list when no
|
|
captions are found.
|
|
"""
|
|
matches = list(_FIGURE_BOUNDARY_RE.finditer(text))
|
|
if not matches:
|
|
return [text]
|
|
segments: list[str] = []
|
|
last = 0
|
|
for m in matches:
|
|
if m.start() > last:
|
|
segments.append(text[last : m.start()])
|
|
last = m.start()
|
|
segments.append(text[last:])
|
|
return [s for s in segments if s.strip()]
|
|
|
|
|
|
@dataclass(frozen = True)
|
|
class Chunk:
|
|
text: str
|
|
token_count: int
|
|
page_number: int | None = None
|
|
source_page_index: int | None = None
|
|
page_char_start: int | None = None
|
|
page_char_end: int | None = None
|
|
line_start: int | None = None
|
|
line_end: int | None = None
|
|
|
|
|
|
TokenCounter = Callable[[str], int]
|
|
|
|
|
|
def _char_token_estimate(text: str) -> int:
|
|
return max(1, (len(text) + 3) // 4)
|
|
|
|
|
|
def _split_on(text: str, separator: str) -> list[str]:
|
|
if separator == "":
|
|
return list(text)
|
|
parts = text.split(separator)
|
|
if len(parts) == 1:
|
|
return parts
|
|
glued: list[str] = []
|
|
for i, part in enumerate(parts):
|
|
if i < len(parts) - 1:
|
|
glued.append(part + separator)
|
|
else:
|
|
if part:
|
|
glued.append(part)
|
|
return [p for p in glued if p]
|
|
|
|
|
|
def _atomic_split(
|
|
text: str,
|
|
separators: tuple[str, ...],
|
|
max_tokens: int,
|
|
count: TokenCounter,
|
|
) -> list[str]:
|
|
if count(text) <= max_tokens:
|
|
return [text]
|
|
for sep in separators:
|
|
pieces = _split_on(text, sep)
|
|
if len(pieces) <= 1:
|
|
continue
|
|
out: list[str] = []
|
|
for piece in pieces:
|
|
if count(piece) <= max_tokens:
|
|
out.append(piece)
|
|
else:
|
|
tail = separators[separators.index(sep) + 1 :]
|
|
out.extend(_atomic_split(piece, tail, max_tokens, count))
|
|
return out
|
|
approx_chars = max(1, max_tokens * 4)
|
|
return [text[i : i + approx_chars] for i in range(0, len(text), approx_chars)]
|
|
|
|
|
|
def _merge(
|
|
pieces: list[str],
|
|
max_tokens: int,
|
|
overlap_tokens: int,
|
|
count: TokenCounter,
|
|
) -> list[str]:
|
|
"""Greedy-merge into <= max_tokens chunks with overlap."""
|
|
chunks: list[str] = []
|
|
buffer: list[str] = []
|
|
buffer_tokens = 0
|
|
for piece in pieces:
|
|
piece_tokens = count(piece)
|
|
if buffer and buffer_tokens + piece_tokens > max_tokens:
|
|
chunks.append("".join(buffer))
|
|
if overlap_tokens > 0:
|
|
overlap: list[str] = []
|
|
running = 0
|
|
for prev in reversed(buffer):
|
|
prev_tokens = count(prev)
|
|
if running + prev_tokens > overlap_tokens:
|
|
break
|
|
overlap.insert(0, prev)
|
|
running += prev_tokens
|
|
buffer = list(overlap)
|
|
buffer_tokens = running
|
|
else:
|
|
buffer = []
|
|
buffer_tokens = 0
|
|
buffer.append(piece)
|
|
buffer_tokens += piece_tokens
|
|
if buffer:
|
|
chunks.append("".join(buffer))
|
|
return [c.strip() for c in chunks if c.strip()]
|
|
|
|
|
|
def _line_bounds(text: str, start: int, end: int) -> tuple[int, int]:
|
|
"""Return 1-based inclusive line numbers for a page-local span."""
|
|
line_start = text.count("\n", 0, start) + 1
|
|
line_end = text.count("\n", 0, max(start, end - 1)) + 1
|
|
return line_start, line_end
|
|
|
|
|
|
def _locate_piece(
|
|
page_text: str,
|
|
piece: str,
|
|
search_cursor: int,
|
|
) -> tuple[int | None, int | None, int | None, int | None, int]:
|
|
idx = page_text.find(piece, search_cursor)
|
|
if idx < 0:
|
|
idx = page_text.find(piece)
|
|
if idx < 0:
|
|
return None, None, None, None, search_cursor
|
|
end = idx + len(piece)
|
|
line_start, line_end = _line_bounds(page_text, idx, end)
|
|
return idx, end, line_start, line_end, idx + 1
|
|
|
|
|
|
# Markdown headings first so layout-aware parser output splits at sections.
|
|
DEFAULT_SEPARATORS: tuple[str, ...] = (
|
|
"\n# ",
|
|
"\n## ",
|
|
"\n### ",
|
|
"\n#### ",
|
|
"\n\n",
|
|
"\n",
|
|
". ",
|
|
" ",
|
|
"",
|
|
)
|
|
|
|
|
|
def chunk_pages(
|
|
pages: list[ParsedPage],
|
|
*,
|
|
max_tokens: int,
|
|
overlap_tokens: int,
|
|
token_counter: TokenCounter | None = None,
|
|
separators: tuple[str, ...] = DEFAULT_SEPARATORS,
|
|
) -> list[Chunk]:
|
|
"""Split pages independently so page_number stays attached to chunks."""
|
|
count = token_counter or _char_token_estimate
|
|
out: list[Chunk] = []
|
|
for page_index, page in enumerate(pages):
|
|
search_cursor = 0
|
|
for segment in _split_at_figure_boundaries(page.text):
|
|
atomic = _atomic_split(segment, separators, max_tokens, count)
|
|
merged = _merge(atomic, max_tokens, overlap_tokens, count)
|
|
for piece in merged:
|
|
start, end, line_start, line_end, search_cursor = _locate_piece(
|
|
page.text,
|
|
piece,
|
|
search_cursor,
|
|
)
|
|
out.append(
|
|
Chunk(
|
|
text = piece,
|
|
token_count = count(piece),
|
|
page_number = page.page_number,
|
|
source_page_index = page_index,
|
|
page_char_start = start,
|
|
page_char_end = end,
|
|
line_start = line_start,
|
|
line_end = line_end,
|
|
)
|
|
)
|
|
return out
|
|
|
|
|
|
_PAGE_SEPARATOR = "\n\n"
|
|
|
|
|
|
def chunk_pages_with_spans(
|
|
pages: list[ParsedPage],
|
|
*,
|
|
max_tokens: int,
|
|
overlap_tokens: int,
|
|
token_counter: TokenCounter | None = None,
|
|
separators: tuple[str, ...] = DEFAULT_SEPARATORS,
|
|
) -> tuple[str, list[Chunk], list[tuple[int, int]]]:
|
|
"""Late-chunking variant: joins pages so the embedder sees the whole doc.
|
|
|
|
Returns ``(full_doc, chunks, char_spans)``; ``char_spans[i]`` is the
|
|
(start, end) char offset of ``chunks[i].text`` inside ``full_doc``.
|
|
Page numbers are recovered by overlap with the original page ranges.
|
|
"""
|
|
count = token_counter or _char_token_estimate
|
|
|
|
parts: list[str] = []
|
|
page_ranges: list[tuple[int, int, int, int | None]] = []
|
|
cursor = 0
|
|
for index, page in enumerate(pages):
|
|
parts.append(page.text)
|
|
start = cursor
|
|
end = cursor + len(page.text)
|
|
page_ranges.append((start, end, index, page.page_number))
|
|
cursor = end
|
|
if index < len(pages) - 1:
|
|
cursor += len(_PAGE_SEPARATOR)
|
|
full_doc = _PAGE_SEPARATOR.join(parts)
|
|
|
|
atomic: list[str] = []
|
|
for segment in _split_at_figure_boundaries(full_doc):
|
|
atomic.extend(_atomic_split(segment, separators, max_tokens, count))
|
|
merged = _merge(atomic, max_tokens, overlap_tokens, count)
|
|
|
|
chunks: list[Chunk] = []
|
|
char_spans: list[tuple[int, int]] = []
|
|
search_cursor = 0
|
|
for piece in merged:
|
|
text = piece.strip()
|
|
if not text:
|
|
continue
|
|
idx = full_doc.find(text, search_cursor)
|
|
if idx < 0:
|
|
# Overlap can push past a chunk's true start; restart from head.
|
|
idx = full_doc.find(text)
|
|
if idx < 0:
|
|
continue
|
|
end_idx = idx + len(text)
|
|
page_locator = _page_for_span(idx, end_idx, page_ranges)
|
|
source_page_index: int | None = None
|
|
page_number: int | None = None
|
|
page_char_start: int | None = None
|
|
page_char_end: int | None = None
|
|
line_start: int | None = None
|
|
line_end: int | None = None
|
|
if page_locator is not None:
|
|
page_start, page_end, page_idx, page_no = page_locator
|
|
source_page_index = page_idx
|
|
page_number = page_no
|
|
page_char_start = max(0, idx - page_start)
|
|
page_char_end = min(page_end, end_idx) - page_start
|
|
line_start, line_end = _line_bounds(
|
|
pages[page_idx].text,
|
|
page_char_start,
|
|
page_char_end,
|
|
)
|
|
chunks.append(
|
|
Chunk(
|
|
text = text,
|
|
token_count = count(text),
|
|
page_number = page_number,
|
|
source_page_index = source_page_index,
|
|
page_char_start = page_char_start,
|
|
page_char_end = page_char_end,
|
|
line_start = line_start,
|
|
line_end = line_end,
|
|
)
|
|
)
|
|
char_spans.append((idx, end_idx))
|
|
# Advance past start (not end) so overlapping next chunk is findable.
|
|
search_cursor = idx + 1
|
|
|
|
return full_doc, chunks, char_spans
|
|
|
|
|
|
def _page_for_span(
|
|
start: int,
|
|
end: int,
|
|
page_ranges: list[tuple[int, int, int, int | None]],
|
|
) -> tuple[int, int, int, int | None] | None:
|
|
for ps, pe, page_index, page_number in page_ranges:
|
|
if start < pe and end > ps:
|
|
return ps, pe, page_index, page_number
|
|
return None
|