Studio RAG: remove late chunking and the chunking_strategy field/selector (single fixed chunker)
This commit is contained in:
Roland Tannous
parent
f955738075
commit
a41fae78eb
16 changed files with 69 additions and 847 deletions
|
|
@ -1,113 +0,0 @@
|
|||
"""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):
|
||||
# Chunk text must equal the full_doc slice 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}
|
||||
# Each page contributes 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
|
||||
# Pages 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: ~80MB, 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: 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,)
|
||||
|
|
@ -7,7 +7,6 @@ returns images when asked, route accepts the mode field, constraint
|
|||
validator rejects illegal combos) run in every test invocation.
|
||||
"""
|
||||
|
||||
import importlib.util
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
|
@ -19,27 +18,6 @@ if str(STUDIO_BACKEND) not in sys.path:
|
|||
sys.path.insert(0, str(STUDIO_BACKEND))
|
||||
|
||||
|
||||
def _rag_route():
|
||||
"""Load ``routes/rag.py`` directly, bypassing the ``routes`` package.
|
||||
|
||||
``from routes.rag import X`` first runs ``routes/__init__.py``, which eagerly
|
||||
imports every router — including the datasets router, whose chain does
|
||||
``from datasets import IterableDataset`` at import time. On a GPU-less CI
|
||||
runner the unsloth bootstrap can leave ``datasets`` half-initialized, so that
|
||||
eager import raises. These tests only need pure helpers from rag.py, so load
|
||||
the file on its own (it has no intra-``routes`` imports).
|
||||
"""
|
||||
mod = sys.modules.get("_rag_route_under_test")
|
||||
if mod is None:
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
"_rag_route_under_test", STUDIO_BACKEND / "routes" / "rag.py"
|
||||
)
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(mod)
|
||||
sys.modules["_rag_route_under_test"] = mod
|
||||
return mod
|
||||
|
||||
|
||||
def test_html_parser_returns_images_when_requested(tmp_path):
|
||||
pytest.importorskip("bs4")
|
||||
pytest.importorskip("lxml")
|
||||
|
|
@ -74,32 +52,14 @@ def test_html_parser_returns_images_when_requested(tmp_path):
|
|||
assert img.nearest_caption == "A tiny figure"
|
||||
|
||||
|
||||
def test_multimodal_late_combo_validator():
|
||||
from fastapi import HTTPException
|
||||
|
||||
_validate_mode_combo = _rag_route()._validate_mode_combo
|
||||
|
||||
# Allowed combos → None.
|
||||
assert _validate_mode_combo("text", "standard") is None
|
||||
assert _validate_mode_combo("text", "late") is None
|
||||
assert _validate_mode_combo("multimodal", "standard") is None
|
||||
|
||||
# Forbidden combo → 400.
|
||||
with pytest.raises(HTTPException) as excinfo:
|
||||
_validate_mode_combo("multimodal", "late")
|
||||
assert excinfo.value.status_code == 400
|
||||
|
||||
|
||||
def test_rag_embedder_matrix_excludes_multimodal_late():
|
||||
def test_rag_embedder_matrix_is_keyed_by_mode():
|
||||
from utils.rag.config import RAG_EMBEDDER_MATRIX, resolve_embedder
|
||||
|
||||
assert ("multimodal", "late") not in RAG_EMBEDDER_MATRIX
|
||||
assert ("text", "standard") in RAG_EMBEDDER_MATRIX
|
||||
assert ("text", "late") in RAG_EMBEDDER_MATRIX
|
||||
assert ("multimodal", "standard") in RAG_EMBEDDER_MATRIX
|
||||
assert "text" in RAG_EMBEDDER_MATRIX
|
||||
assert "multimodal" in RAG_EMBEDDER_MATRIX
|
||||
|
||||
# Unknown combos fall back to the legacy default, not KeyError.
|
||||
fallback = resolve_embedder("multimodal", "late")
|
||||
# Unknown modes fall back to the default, not KeyError.
|
||||
fallback = resolve_embedder("unknown-mode")
|
||||
assert isinstance(fallback, str) and fallback
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -93,7 +93,6 @@ def test_multimodal_subprocess_emits_image_and_caption_chunks(
|
|||
overlap = 20,
|
||||
batch_size = 4,
|
||||
out_queue = out_queue,
|
||||
chunking_strategy = "standard",
|
||||
mode = "multimodal",
|
||||
document_id = "test-doc-1",
|
||||
)
|
||||
|
|
|
|||
|
|
@ -2,8 +2,7 @@
|
|||
|
||||
Full end-to-end reingest needs a running studio + a real embedder; that's
|
||||
covered manually via the curl smoke flow in the plan. Here we cover the
|
||||
parts that are testable without external models: payload validation and
|
||||
the (multimodal, late) constraint propagation.
|
||||
parts that are testable without external models: payload validation.
|
||||
"""
|
||||
|
||||
import importlib.util
|
||||
|
|
@ -43,22 +42,12 @@ def test_reingest_request_accepts_all_optional_fields():
|
|||
ReingestKBRequest = _rag_route().ReingestKBRequest
|
||||
|
||||
empty = ReingestKBRequest()
|
||||
assert empty.chunking_strategy is None
|
||||
assert empty.mode is None
|
||||
assert empty.embedding_model is None
|
||||
|
||||
partial = ReingestKBRequest(chunking_strategy = "late")
|
||||
assert partial.chunking_strategy == "late"
|
||||
assert partial.mode is None
|
||||
|
||||
|
||||
def test_reingest_request_rejects_unknown_strategy():
|
||||
from pydantic import ValidationError
|
||||
|
||||
ReingestKBRequest = _rag_route().ReingestKBRequest
|
||||
|
||||
with pytest.raises(ValidationError):
|
||||
ReingestKBRequest(chunking_strategy = "telekinetic")
|
||||
partial = ReingestKBRequest(mode = "multimodal")
|
||||
assert partial.mode == "multimodal"
|
||||
assert partial.embedding_model is None
|
||||
|
||||
|
||||
def test_reingest_request_rejects_unknown_mode():
|
||||
|
|
@ -68,14 +57,3 @@ def test_reingest_request_rejects_unknown_mode():
|
|||
|
||||
with pytest.raises(ValidationError):
|
||||
ReingestKBRequest(mode = "augmented")
|
||||
|
||||
|
||||
def test_constraint_still_enforced_for_reingest_combos():
|
||||
"""The combination guard is shared with create — verify it still bites."""
|
||||
from fastapi import HTTPException
|
||||
|
||||
_validate_mode_combo = _rag_route()._validate_mode_combo
|
||||
|
||||
with pytest.raises(HTTPException) as excinfo:
|
||||
_validate_mode_combo("multimodal", "late")
|
||||
assert excinfo.value.status_code == 400
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue