ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source 
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
Download patch file Download diff file Expand all files Collapse all files

127
studio/backend/core/rag/chunking.py
View file

@ -100,28 +100,31 @@ def _merge(
return [c.strip() for c in chunks if c.strip()]
DEFAULT_SEPARATORS: tuple[str, ...] = (
# Markdown heading boundaries first — when the parser emits
# layout-aware Markdown (PDF via pymupdf4llm, DOCX via mammoth,
# HTML via markdownify) chunks split at section breaks rather
# than mid-paragraph. Falls back to the original separators on
# plain text input where headings are absent.
"\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, ...] = (
# Markdown heading boundaries first — when the parser emits
# layout-aware Markdown (PDF via pymupdf4llm, DOCX via mammoth,
# HTML via markdownify) chunks split at section breaks rather
# than mid-paragraph. Falls back to the original separators on
# plain text input where headings are absent.
"\n# ",
"\n## ",
"\n### ",
"\n#### ",
"\n\n",
"\n",
". ",
" ",
"",
),
separators: tuple[str, ...] = DEFAULT_SEPARATORS,
) -> list[Chunk]:
"""Split parsed pages into overlapping chunks.
@ -142,3 +145,95 @@ def chunk_pages(
)
)
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-friendly variant of :func:`chunk_pages`.
Joins all pages into a single document so the embedder sees the
whole text in one pass (that's the point of late chunking — chunk
vectors that carry full-document context via the model's
bidirectional attention).
Returns ``(full_doc, chunks, char_spans)`` where
``char_spans[i] = (start, end)`` are byte-character offsets of
``chunks[i].text`` inside ``full_doc``. The embedder layer maps
char spans → token spans via the tokenizer's offsets_mapping and
mean-pools per chunk.
Page-number metadata on each :class:`Chunk` is recovered from the
chunk's char span — the first page whose range overlaps the chunk
wins. PDFs keep useful citations even though chunking ignores page
boundaries here.
"""
count = token_counter or _char_token_estimate
parts: list[str] = []
page_ranges: list[tuple[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, page.page_number))
cursor = end
if index < len(pages) - 1:
cursor += len(_PAGE_SEPARATOR)
full_doc = _PAGE_SEPARATOR.join(parts)
atomic = _atomic_split(full_doc, 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 the search cursor past a chunk's true
# start — restart from the document head as a fallback.
idx = full_doc.find(text)
if idx < 0:
# Chunker output diverged from the source (rare — happens
# if a separator-splice mangled the text). Skip the chunk
# rather than corrupt the vector store with a wrong span.
continue
end_idx = idx + len(text)
page_number = _page_for_span(idx, end_idx, page_ranges)
chunks.append(
Chunk(
text = text,
token_count = count(text),
page_number = page_number,
)
)
char_spans.append((idx, end_idx))
# Advance past the *start* of this chunk so an overlapping
# next chunk can still be found.
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 | None]],
) -> int | None:
for ps, pe, pn in page_ranges:
if start < pe and end > ps:
return pn
return None

238
studio/backend/core/rag/embeddings.py
View file

@ -100,3 +100,241 @@ def token_counter(model_name: str | None = None):
return max(1, len(text) // 4)
return _count
# ------------------------------------------------------------------
# Late chunking (Phase 3B-late)
# ------------------------------------------------------------------
_LATE_WINDOW_OVERLAP_TOKENS = 512
def late_chunk_encode(
doc_text: str,
char_spans: list[tuple[int, int]],
*,
model_name: str | None = None,
normalize: bool = True,
):
"""Embed each chunk via late-chunking pooling.
Single forward pass over the full document, then mean-pool the
token embeddings whose offset ranges fall inside each chunk's
char span. Chunks therefore carry full-document context via the
encoder's bidirectional attention — Jina's published technique,
works with any encoder that exposes per-token outputs.
When the doc exceeds the embedder's context, falls back to
windowed late chunking with a 512-token overlap between windows
so cross-window context is partially preserved.
"""
import numpy as np
if not char_spans:
return []
model = get_embedder(model_name)
tokenizer = model.tokenizer
max_length = int(getattr(model, "max_seq_length", None) or 8192)
encoded = tokenizer(
doc_text,
return_tensors = "pt",
return_offsets_mapping = True,
add_special_tokens = True,
truncation = False,
)
offsets = encoded.pop("offset_mapping")[0].tolist()
n_tokens = int(encoded["input_ids"].shape[1])
if n_tokens <= max_length:
token_embeddings = _encode_tokens(model, encoded)
return _pool_spans(
token_embeddings,
offsets,
char_spans,
normalize = normalize,
np_module = np,
model = model,
doc_text = doc_text,
)
logger.info(
"Late chunking: doc has %d tokens > model max %d; using windowed pass",
n_tokens,
max_length,
)
return _windowed_late_chunk_encode(
doc_text = doc_text,
char_spans = char_spans,
model = model,
max_length = max_length,
normalize = normalize,
np_module = np,
)
def _encode_tokens(model, encoded):
"""Run the embedder's underlying transformer to get per-token last_hidden_state."""
import torch
transformer = model[0].auto_model
device = next(transformer.parameters()).device
inputs_on_device = {k: v.to(device) for k, v in encoded.items()}
with torch.no_grad():
outputs = transformer(**inputs_on_device)
return outputs.last_hidden_state[0].detach().cpu().numpy()
def _pool_spans(
token_embeddings,
offsets,
char_spans,
*,
normalize: bool,
np_module,
model,
doc_text: str,
token_index_offset: int = 0,
):
"""Mean-pool token embeddings per (char_start, char_end) span.
`token_index_offset` shifts char_span-derived token indices into
a sub-window's local frame (used by the windowed code path).
"""
vectors = []
n_rows = token_embeddings.shape[0]
for char_start, char_end in char_spans:
# Special tokens (CLS / SEP) report offsets (0, 0) — exclude them.
indices = [
i - token_index_offset
for i, (ts, te) in enumerate(offsets)
if te > ts and te > char_start and ts < char_end
]
indices = [i for i in indices if 0 <= i < n_rows]
if not indices:
# Fall back to a standalone encode of the chunk text — rare
# (would mean tokenizer produced zero non-special tokens for
# the span), but keeps the pipeline alive.
vec = model.encode(
doc_text[char_start:char_end],
normalize_embeddings = normalize,
convert_to_numpy = True,
show_progress_bar = False,
)
vectors.append(vec)
continue
pooled = token_embeddings[indices].mean(axis = 0)
if normalize:
denom = float(np_module.linalg.norm(pooled))
if denom > 0:
pooled = pooled / denom
vectors.append(pooled)
return vectors
def _windowed_late_chunk_encode(
*,
doc_text: str,
char_spans: list[tuple[int, int]],
model,
max_length: int,
normalize: bool,
np_module,
):
"""Doc exceeds context window — slice into overlapping windows.
Each chunk is pooled against the window that contains the most of
its tokens. The 512-token window overlap means chunks near a
boundary still see context from both sides.
"""
import torch
tokenizer = model.tokenizer
transformer = model[0].auto_model
device = next(transformer.parameters()).device
full = tokenizer(
doc_text,
return_tensors = "pt",
return_offsets_mapping = True,
add_special_tokens = False,
truncation = False,
)
all_input_ids = full["input_ids"][0]
all_offsets = full["offset_mapping"][0].tolist()
n_tokens = int(all_input_ids.shape[0])
stride = max(1, max_length - _LATE_WINDOW_OVERLAP_TOKENS)
# Build (start_token, end_token) windows.
windows: list[tuple[int, int]] = []
pos = 0
while pos < n_tokens:
end = min(pos + max_length, n_tokens)
windows.append((pos, end))
if end >= n_tokens:
break
pos += stride
# Cache window → token embeddings (only encode when needed).
window_embeddings: dict[int, "np_module.ndarray"] = {}
def _window_embeddings(window_index: int):
if window_index in window_embeddings:
return window_embeddings[window_index]
ws, we = windows[window_index]
win_ids = all_input_ids[ws:we].unsqueeze(0).to(device)
win_attn = torch.ones_like(win_ids)
with torch.no_grad():
outputs = transformer(input_ids = win_ids, attention_mask = win_attn)
emb = outputs.last_hidden_state[0].detach().cpu().numpy()
window_embeddings[window_index] = emb
return emb
vectors = []
for char_start, char_end in char_spans:
# Collect global token indices in the chunk.
chunk_token_indices = [
i
for i, (ts, te) in enumerate(all_offsets)
if te > ts and te > char_start and ts < char_end
]
if not chunk_token_indices:
vec = model.encode(
doc_text[char_start:char_end],
normalize_embeddings = normalize,
convert_to_numpy = True,
show_progress_bar = False,
)
vectors.append(vec)
continue
# Pick the window covering the most of this chunk's tokens.
best_window = 0
best_overlap = 0
for wi, (ws, we) in enumerate(windows):
overlap = sum(1 for ti in chunk_token_indices if ws <= ti < we)
if overlap > best_overlap:
best_overlap = overlap
best_window = wi
ws, _we = windows[best_window]
emb = _window_embeddings(best_window)
local_indices = [
ti - ws
for ti in chunk_token_indices
if ws <= ti < ws + emb.shape[0]
]
if not local_indices:
vec = model.encode(
doc_text[char_start:char_end],
normalize_embeddings = normalize,
convert_to_numpy = True,
show_progress_bar = False,
)
vectors.append(vec)
continue
pooled = emb[local_indices].mean(axis = 0)
if normalize:
denom = float(np_module.linalg.norm(pooled))
if denom > 0:
pooled = pooled / denom
vectors.append(pooled)
return vectors

202
studio/backend/core/rag/ingestion.py
View file

@ -54,74 +54,172 @@ def _subprocess_worker(
overlap: int,
batch_size: int,
out_queue: Any,
chunking_strategy: str = "standard",
mode: str = "text",
) -> None:
try:
from core.rag.chunking import chunk_pages
from core.rag.chunking import chunk_pages, chunk_pages_with_spans
from core.rag.parsers import parse
out_queue.put({"type": "progress", "stage": "parse", "progress": 0.05})
# want_images stays False for the text-only ingestion path; the
# multimodal path (Phase 3B-multimodal) will flip this based on
# the KB's mode.
parsed = parse(Path(stored_path), want_images = False)
# want_images is True only for multimodal KBs. The image side of
# the pipeline lands in Phase 3B-multimodal; for now the parser
# collects the bytes anyway in case we want them later, but only
# the text pages are consumed.
parsed = parse(Path(stored_path), want_images = (mode == "multimodal"))
pages = parsed.pages
if not pages:
out_queue.put({"type": "error", "error": "no extractable text in document"})
return
out_queue.put({"type": "progress", "stage": "load_model", "progress": 0.1})
from core.rag.embeddings import get_embedder, token_counter
from core.rag.embeddings import (
get_embedder,
late_chunk_encode,
token_counter,
)
model = get_embedder(model_name)
counter = token_counter(model_name)
dim = int(model.get_sentence_embedding_dimension())
out_queue.put({"type": "dim", "dim": dim})
out_queue.put({"type": "progress", "stage": "chunk", "progress": 0.2})
chunks = chunk_pages(
pages,
max_tokens = chunk_size,
overlap_tokens = overlap,
token_counter = counter,
)
if not chunks:
out_queue.put({"type": "error", "error": "chunker produced no chunks"})
return
total = len(chunks)
for i in range(0, total, batch_size):
batch = chunks[i : i + batch_size]
vectors = model.encode(
[c.text for c in batch],
if chunking_strategy == "late":
_run_late_chunking(
pages = pages,
chunk_size = chunk_size,
overlap = overlap,
counter = counter,
model_name = model_name,
late_chunk_encode = late_chunk_encode,
out_queue = out_queue,
)
else:
_run_standard_chunking(
pages = pages,
chunk_size = chunk_size,
overlap = overlap,
counter = counter,
batch_size = batch_size,
normalize_embeddings = True,
convert_to_numpy = True,
show_progress_bar = False,
model = model,
chunk_pages = chunk_pages,
out_queue = out_queue,
)
out_queue.put(
{
"type": "chunks_batch",
"first_index": i,
"chunks": [
{
"text": c.text,
"token_count": c.token_count,
"page_number": c.page_number,
}
for c in batch
],
"vectors": vectors.tolist(),
}
)
progress = 0.3 + 0.65 * min(1.0, (i + len(batch)) / total)
out_queue.put({"type": "progress", "stage": "embed", "progress": progress})
out_queue.put({"type": "complete", "num_chunks": total})
except Exception as exc: # noqa: BLE001
logger.exception("ingestion subprocess failed")
out_queue.put({"type": "error", "error": f"{type(exc).__name__}: {exc}"})
def _run_standard_chunking(
*,
pages,
chunk_size,
overlap,
counter,
batch_size,
model,
chunk_pages,
out_queue,
) -> None:
out_queue.put({"type": "progress", "stage": "chunk", "progress": 0.2})
chunks = chunk_pages(
pages,
max_tokens = chunk_size,
overlap_tokens = overlap,
token_counter = counter,
)
if not chunks:
out_queue.put({"type": "error", "error": "chunker produced no chunks"})
return
total = len(chunks)
for i in range(0, total, batch_size):
batch = chunks[i : i + batch_size]
vectors = model.encode(
[c.text for c in batch],
batch_size = batch_size,
normalize_embeddings = True,
convert_to_numpy = True,
show_progress_bar = False,
)
out_queue.put(
{
"type": "chunks_batch",
"first_index": i,
"chunks": [
{
"text": c.text,
"token_count": c.token_count,
"page_number": c.page_number,
}
for c in batch
],
"vectors": vectors.tolist(),
}
)
progress = 0.3 + 0.65 * min(1.0, (i + len(batch)) / total)
out_queue.put({"type": "progress", "stage": "embed", "progress": progress})
out_queue.put({"type": "complete", "num_chunks": total})
def _run_late_chunking(
*,
pages,
chunk_size,
overlap,
counter,
model_name,
late_chunk_encode,
out_queue,
) -> None:
"""Late chunking: chunk once over the whole doc, embed in a single pass.
There's no per-batch streaming here — the whole doc is encoded in
one forward pass (or one per window for long docs). We ship all
chunks back to the parent in one message; the parent's pump still
handles them via the same chunks_batch handler.
"""
from core.rag.chunking import chunk_pages_with_spans
out_queue.put({"type": "progress", "stage": "chunk", "progress": 0.2})
full_doc, chunks, char_spans = chunk_pages_with_spans(
pages,
max_tokens = chunk_size,
overlap_tokens = overlap,
token_counter = counter,
)
if not chunks:
out_queue.put({"type": "error", "error": "chunker produced no chunks"})
return
out_queue.put({"type": "progress", "stage": "embed", "progress": 0.4})
vectors = late_chunk_encode(
full_doc,
char_spans,
model_name = model_name,
normalize = True,
)
out_queue.put({"type": "progress", "stage": "embed", "progress": 0.9})
out_queue.put(
{
"type": "chunks_batch",
"first_index": 0,
"chunks": [
{
"text": c.text,
"token_count": c.token_count,
"page_number": c.page_number,
}
for c in chunks
],
"vectors": [v.tolist() for v in vectors],
}
)
out_queue.put({"type": "complete", "num_chunks": len(chunks)})
# ------------------------------------------------------------------
# Job manager (parent side)
# ------------------------------------------------------------------
@ -390,14 +488,26 @@ def enqueue_ingestion(
kb_id: str | None = None,
thread_id: str | None = None,
embedding_model: str | None = None,
chunking_strategy: str = "standard",
mode: str = "text",
) -> str:
"""Create the job row, spawn the subprocess, and start the pump thread.
Returns the job_id. The caller can poll via ``GET /api/rag/jobs/{job_id}/events``
or read the ``rag_ingestion_jobs`` table directly.
chunking_strategy / mode default to today's behaviour. KB-scoped
uploads should pass the KB's stored values; per-thread uploads
default unless an override is set in chat_settings.
"""
from utils.rag.config import resolve_embedder
scope = _scope_for(kb_id, thread_id)
model_name = embedding_model or RAG_EMBEDDING_MODEL
model_name = (
embedding_model
or resolve_embedder(mode, chunking_strategy)
or RAG_EMBEDDING_MODEL
)
job_id = str(uuid4())
with get_connection() as conn:
conn.execute(
@ -424,6 +534,8 @@ def enqueue_ingestion(
RAG_CHUNK_OVERLAP,
RAG_EMBED_BATCH_SIZE,
out_queue,
chunking_strategy,
mode,
),
daemon = True,
)

16
studio/backend/routes/rag.py
View file

@ -289,6 +289,8 @@ def _start_ingestion(
kb_id: str | None,
thread_id: str | None,
embedding_model: str,
chunking_strategy: str = "standard",
mode: str = "text",
) -> UploadResponse:
document_id = str(uuid4())
with get_connection() as conn:
@ -317,6 +319,8 @@ def _start_ingestion(
kb_id = kb_id,
thread_id = thread_id,
embedding_model = embedding_model,
chunking_strategy = chunking_strategy,
mode = mode,
)
return UploadResponse(document_id = document_id, job_id = job_id, filename = filename)
@ -432,6 +436,16 @@ async def upload_kb_document(
) -> UploadResponse:
kb_row = _kb_or_404(kb_id)
stored_path, filename, byte_size = await _save_upload(file)
# Defensive .get() — rows fetched through a connection that pre-dates
# the Phase 3 schema (e.g. in tests) lack chunking_strategy/mode;
# fall back to the same defaults as the column.
kb_keys = kb_row.keys() if hasattr(kb_row, "keys") else ()
chunking_strategy = (
kb_row["chunking_strategy"]
if "chunking_strategy" in kb_keys
else "standard"
)
mode = kb_row["mode"] if "mode" in kb_keys else "text"
return _start_ingestion(
filename = filename,
stored_path = stored_path,
@ -440,6 +454,8 @@ async def upload_kb_document(
kb_id = kb_id,
thread_id = None,
embedding_model = kb_row["embedding_model"],
chunking_strategy = chunking_strategy,
mode = mode,
)

13
studio/frontend/src/features/chat/chat-settings-sheet.tsx
View file

@ -1232,6 +1232,7 @@ export function ChatSettingsPanel({
) : null}
{knowledgeBases.map((kb) => {
const isActive = kb.id === activeKbId;
const isLate = kb.chunking_strategy === "late";
return (
<DropdownMenuItem
key={kb.id}
@ -1243,7 +1244,17 @@ export function ChatSettingsPanel({
setRagSource({ kind: "kb", kbId: kb.id })
}
>
<span className="truncate">{kb.name}</span>
<span className="flex min-w-0 items-center gap-1.5 truncate">
<span className="truncate">{kb.name}</span>
{isLate ? (
<span
className="rounded-sm bg-amber-500/15 px-1 text-[10px] font-medium text-amber-700 dark:text-amber-300"
title="Late chunking enabled"
>
⚡ Late
</span>
) : null}
</span>
<button
type="button"
aria-label={`Delete ${kb.name}`}

45
studio/frontend/src/features/rag/components/kb-create-dialog.tsx
View file

@ -12,8 +12,15 @@ import {
} from "@/components/ui/dialog";
import { Input } from "@/components/ui/input";
import { Label } from "@/components/ui/label";
import {
Select,
SelectContent,
SelectItem,
SelectTrigger,
SelectValue,
} from "@/components/ui/select";
import { useState } from "react";
import type { KnowledgeBase } from "../api/rag-api";
import type { ChunkingStrategy, KnowledgeBase } from "../api/rag-api";
import { useKnowledgeBases } from "../hooks/use-knowledge-bases";
export function KBCreateDialog({
@ -29,6 +36,8 @@ export function KBCreateDialog({
const [name, setName] = useState("");
const [description, setDescription] = useState("");
const [embeddingModel, setEmbeddingModel] = useState("");
const [chunkingStrategy, setChunkingStrategy] =
useState<ChunkingStrategy>("standard");
const [submitting, setSubmitting] = useState(false);
const [error, setError] = useState<string | null>(null);
@ -36,6 +45,7 @@ export function KBCreateDialog({
setName("");
setDescription("");
setEmbeddingModel("");
setChunkingStrategy("standard");
setError(null);
setSubmitting(false);
};
@ -50,6 +60,7 @@ export function KBCreateDialog({
name: name.trim(),
description: description.trim() || undefined,
embedding_model: embeddingModel.trim() || undefined,
chunking_strategy: chunkingStrategy,
});
onCreated?.(kb);
reset();
@ -98,13 +109,43 @@ export function KBCreateDialog({
placeholder="What's in this KB?"
/>
</div>
<div className="flex flex-col gap-2">
<Label htmlFor="kb-strategy">Chunking strategy</Label>
<Select
value={chunkingStrategy}
onValueChange={(v) => setChunkingStrategy(v as ChunkingStrategy)}
>
<SelectTrigger id="kb-strategy">
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectItem value="standard">
Standard — heading-aware recursive splitter
</SelectItem>
<SelectItem value="late">
Late chunking — single-pass embedder, slower ingest
</SelectItem>
</SelectContent>
</Select>
<p className="text-[11px] text-muted-foreground">
Late chunking embeds the whole document in one pass, so each
chunk vector carries full-document context. Slower to ingest
(one forward pass per doc) but improves retrieval on long,
cross-referenced text. Cannot be combined with multimodal
mode.
</p>
</div>
<div className="flex flex-col gap-2">
<Label htmlFor="kb-model">Embedding model (optional)</Label>
<Input
id="kb-model"
value={embeddingModel}
onChange={(e) => setEmbeddingModel(e.target.value)}
placeholder="Defaults to BAAI/bge-small-en-v1.5"
placeholder={
chunkingStrategy === "late"
? "Defaults to nomic-ai/nomic-embed-text-v1.5"
: "Defaults to BAAI/bge-small-en-v1.5"
}
/>
</div>
{error ? (

12
studio/frontend/src/features/rag/components/kb-list.tsx
View file

@ -58,7 +58,17 @@ export function KBList({
onClick={() => onSelect(kb)}
>
<div className="flex min-w-0 flex-col">
<span className="truncate text-sm font-medium">{kb.name}</span>
<span className="flex min-w-0 items-center gap-1.5 truncate text-sm font-medium">
<span className="truncate">{kb.name}</span>
{kb.chunking_strategy === "late" ? (
<span
className="shrink-0 rounded-sm bg-amber-500/15 px-1 text-[10px] font-medium text-amber-700 dark:text-amber-300"
title="Late chunking enabled"
>
⚡ Late
</span>
) : null}
</span>
{kb.description ? (
<span className="truncate text-xs text-muted-foreground">
{kb.description}

112
tests/python/test_rag_late_chunking.py Normal file
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,)