diff --git a/studio/backend/core/rag/chunking.py b/studio/backend/core/rag/chunking.py index 4178b47c8a..86488a9ef6 100644 --- a/studio/backend/core/rag/chunking.py +++ b/studio/backend/core/rag/chunking.py @@ -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 diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 4953d48ffe..b3ea2de430 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -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 diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index 226fd53129..2d99c7e4d7 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -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, ) diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index d292db4ba6..bbf0b22834 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -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, ) diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index bef9fde24a..5803832120 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -1232,6 +1232,7 @@ export function ChatSettingsPanel({ ) : null} {knowledgeBases.map((kb) => { const isActive = kb.id === activeKbId; + const isLate = kb.chunking_strategy === "late"; return ( - {kb.name} + + {kb.name} + {isLate ? ( + + ⚡ Late + + ) : null} +