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:
parent
4c1ab745d6
commit
673b7f86ba
8 changed files with 700 additions and 65 deletions
112
tests/python/test_rag_late_chunking.py
Normal file
112
tests/python/test_rag_late_chunking.py
Normal 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,)
|
||||
Loading…
Add table
Add a link
Reference in a new issue