From 34d24bcf977f6b9fce7dba6324af690b66ea8ad2 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 31 May 2026 14:18:03 +0000 Subject: [PATCH] Studio RAG trim: remove late-chunking strategy Late chunking (the Jina single-pass technique) is a non-default embedding path (chunking_strategy defaults to 'standard'). It adds a second full-document forward pass with windowed token pooling for a marginal long-document gain that the standard per-chunk encoder already covers at R@5 = 1.0 on the gold set. - remove late_chunk_encode / _windowed_late_chunk_encode / _pool_spans / _encode_tokens from embeddings.py (205 lines) - remove the chunking_strategy == 'late' branch and _run_late_chunking from ingestion.py (75 lines); 'late' now degrades to standard chunking 42 RAG tests pass. --- studio/backend/core/rag/embeddings.py | 205 -------------------------- studio/backend/core/rag/ingestion.py | 75 ---------- 2 files changed, 280 deletions(-) diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 7cf3141ad1..fe4fb9e81f 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -247,208 +247,3 @@ def token_counter(model_name: str | None = None): return _count - -# --- Late chunking (Jina technique) --- - -_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, -): - """Single forward pass over the doc, mean-pool token embeddings per chunk span.""" - 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): - 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.""" - vectors = [] - n_rows = token_embeddings.shape[0] - for char_start, char_end in char_spans: - # Skip special tokens whose offsets are (0, 0). - 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: - 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 > ctx window: pool each chunk against the window containing most of its tokens.""" - 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) - - 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 - - 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: - 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 - 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 ea60277b8b..40990e1371 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -158,7 +158,6 @@ def _subprocess_worker( out_queue.put({"type": "progress", "stage": "load_model", "progress": 0.1}) from core.rag.embeddings import ( get_embedder, - late_chunk_encode, token_counter, ) @@ -167,19 +166,6 @@ def _subprocess_worker( dim = int(model.get_sentence_embedding_dimension()) out_queue.put({"type": "dim", "dim": dim}) - if chunking_strategy == "late": - _run_late_chunking( - pages = pages, - stored_path = Path(stored_path), - chunk_size = chunk_size, - overlap = overlap, - counter = counter, - model_name = model_name, - late_chunk_encode = late_chunk_encode, - out_queue = out_queue, - ) - return - text_count = _run_standard_chunking( pages = pages, stored_path = Path(stored_path), @@ -382,67 +368,6 @@ def _stream_image_chunks( return len(out_chunks) -def _run_late_chunking( - *, - pages, - stored_path, - chunk_size, - overlap, - counter, - model_name, - late_chunk_encode, - out_queue, -) -> None: - """Chunk once, embed in one pass, ship all chunks in one chunks_batch.""" - from core.rag.chunking import chunk_pages_with_spans - from core.rag.locators import pdf_regions_for_chunks - - 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, - ) - pdf_regions = pdf_regions_for_chunks(stored_path, pages, chunks) - - 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, - "source_page_index": c.source_page_index, - "page_char_start": c.page_char_start, - "page_char_end": c.page_char_end, - "line_start": c.line_start, - "line_end": c.line_end, - "pdf_regions": pdf_regions[index], - "kind": "text", - } - for index, c in enumerate(chunks) - ], - "vectors": [v.tolist() for v in vectors], - } - ) - out_queue.put({"type": "complete", "num_chunks": len(chunks)}) - - # --- Job manager (parent side) ---