Studio: late chunking opt-in per KB (Phase 3B-late)

When a KB has chunking_strategy = 'late', ingestion takes a separate
code path that embeds the full document in a single forward pass and
mean-pools token embeddings per chunk span. Each chunk vector carries
full-document context via the encoder's bidirectional attention —
Jina's published technique, ~+6.5 nDCG@10 on long docs.

Backend
- chunking.py: new chunk_pages_with_spans() that joins all pages into a
  single full_doc, runs the existing recursive splitter, and returns
  per-chunk (char_start, char_end) offsets. Page-number metadata is
  recovered by overlap with the original page ranges so PDF citations
  still work. Existing chunk_pages() unchanged.
- embeddings.py: new late_chunk_encode(doc_text, char_spans). Tokenizes
  the doc with return_offsets_mapping, runs the underlying transformer
  to get per-token last_hidden_state, then mean-pools per chunk span.
  When the doc exceeds the embedder's context, falls back to windowed
  late chunking with a 512-token overlap so cross-window context is
  partially preserved.
- ingestion.py _subprocess_worker: branches on chunking_strategy.
  'late' path: chunk_pages_with_spans -> late_chunk_encode -> one big
  chunks_batch message. 'standard' path unchanged. Both reuse the same
  parent-side pump.
- ingestion.enqueue_ingestion: new chunking_strategy + mode kwargs;
  defaults to 'standard' / 'text' for legacy callers. embedder model
  resolved via resolve_embedder() from the (mode, strategy) matrix.
- routes/rag.py: KB-doc upload reads chunking_strategy + mode from the
  KB row (defensive .get for pre-Phase-3 schemas) and threads them
  through _start_ingestion.

Frontend
- kb-create-dialog.tsx: new "Chunking strategy" select with Standard /
  Late options. Embedding-model placeholder switches to nomic when
  Late is picked. createKB request now carries chunking_strategy.
- kb-list.tsx + chat-settings-sheet.tsx: small " Late" badge next to
  late-chunking KB names in the settings KB list and the chat sidebar
  dropdown so users see the mode at a glance.

Tests
- test_rag_late_chunking.py: pure-python tests for chunk_pages_with_spans
  (chunks index back into full_doc; page numbers inherited by overlap;
  pages joined with blank line). A server-marked test loads
  all-MiniLM-L6-v2 to exercise late_chunk_encode end-to-end.

No multimodal yet; that's Phase 3B-multimodal (next PR).
This commit is contained in:
Roland Tannous 2026-05-24 12:18:09 +04:00
commit 673b7f86ba
8 changed files with 700 additions and 65 deletions

View file

@ -0,0 +1,112 @@
"""Late chunking tests (Phase 3B-late).
Pure-python coverage of `chunk_pages_with_spans` runs always. The
encoder test loads a small SentenceTransformer and is gated behind the
existing `server` marker so default `pytest` runs skip it.
"""
import sys
from pathlib import Path
import pytest
REPO_ROOT = Path(__file__).resolve().parents[2]
STUDIO_BACKEND = REPO_ROOT / "studio" / "backend"
if str(STUDIO_BACKEND) not in sys.path:
sys.path.insert(0, str(STUDIO_BACKEND))
from core.rag.chunking import chunk_pages_with_spans
from core.rag.parsers import ParsedPage
def _wc_counter(text: str) -> int:
return max(1, len(text.split()))
def test_spans_index_back_to_full_doc_text():
pages = [
ParsedPage(text = "# Section A\n\n" + ("alpha " * 20), page_number = 1),
ParsedPage(text = "# Section B\n\n" + ("beta " * 20), page_number = 2),
]
full_doc, chunks, char_spans = chunk_pages_with_spans(
pages,
max_tokens = 12,
overlap_tokens = 0,
token_counter = _wc_counter,
)
assert chunks
assert len(chunks) == len(char_spans)
for chunk, (start, end) in zip(chunks, char_spans):
# The chunk text must be exactly the slice of full_doc it claims.
assert full_doc[start:end] == chunk.text
def test_chunks_inherit_page_number_by_overlap():
pages = [
ParsedPage(text = "page-one text here", page_number = 1),
ParsedPage(text = "page-two text here", page_number = 2),
]
_full_doc, chunks, _spans = chunk_pages_with_spans(
pages,
max_tokens = 4,
overlap_tokens = 0,
token_counter = _wc_counter,
)
pages_seen = {c.page_number for c in chunks}
assert pages_seen <= {1, 2}
# Both pages should contribute at least one chunk.
assert 1 in pages_seen
assert 2 in pages_seen
def test_full_doc_joins_pages_with_blank_line_separator():
pages = [
ParsedPage(text = "first", page_number = 1),
ParsedPage(text = "second", page_number = 2),
]
full_doc, _chunks, _spans = chunk_pages_with_spans(
pages,
max_tokens = 5,
overlap_tokens = 0,
token_counter = _wc_counter,
)
assert "first" in full_doc
assert "second" in full_doc
# The two pages must be separated by exactly one blank line.
assert "first\n\nsecond" in full_doc
@pytest.mark.server
def test_late_chunk_encode_returns_one_vector_per_span():
pytest.importorskip("sentence_transformers")
pytest.importorskip("torch")
# all-MiniLM-L6-v2 is ~80MB and embeds at 384 dims.
import os
os.environ.setdefault(
"UNSLOTH_RAG_EMBEDDING_MODEL",
"sentence-transformers/all-MiniLM-L6-v2",
)
from core.rag import embeddings as embeddings_module
embeddings_module._model = None # force re-load
embeddings_module._model_name = None
doc_text = (
"# Intro\n\n"
"The quick brown fox jumps over the lazy dog.\n\n"
"# Methods\n\n"
"We trained the model on a corpus of 100M tokens.\n\n"
"# Results\n\n"
"Accuracy improved by 12% over the baseline."
)
# char_spans for three chunks — one per section, picked manually.
char_spans = [
(doc_text.index("The quick"), doc_text.index("\n\n# Methods")),
(doc_text.index("We trained"), doc_text.index("\n\n# Results")),
(doc_text.index("Accuracy"), len(doc_text)),
]
vectors = embeddings_module.late_chunk_encode(doc_text, char_spans)
assert len(vectors) == len(char_spans)
dim = vectors[0].shape[0]
for v in vectors:
assert v.shape == (dim,)