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.
This commit is contained in:
Daniel Han
parent
b085b37b69
commit
34d24bcf97
2 changed files with 0 additions and 280 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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) ---
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue