92994e8b83
Backend (studio/backend/): - core/rag/: parsers (PDF/TXT/MD/DOCX/HTML via pypdf/python-docx/bs4), recursive token-aware chunker, embeddings singleton via FastSentenceTransformer.from_pretrained(for_inference=True), Qdrant local vector store, bm25s lexical index, RRF hybrid retrieval, spawn-subprocess ingestion job with SSE progress, optional CrossEncoder reranker (off-by-default). - routes/rag.py: KB CRUD, doc upload (KB + per-thread), doc list/delete, ingestion SSE, hybrid+rerank search, thread-index list/clear. - routes/chat_history.py: purge thread RAG artifacts on thread delete and clear-all (rag_documents has no FK cascade to chat_threads so uploads work on un-persisted threads). - studio.db gains 4 RAG tables; storage_roots gains rag_*() helpers. - auth/authentication.py: get_current_subject_sse accepts ?token=... so EventSource can stream ingestion progress. Frontend (studio/frontend/): - features/rag/: api client, Zustand store, hooks, dropzone, KB list, doc rows, ingestion-progress, thread-index list components. - Settings dialog gains a Knowledge Bases tab (master/detail + thread documents list); /knowledge-bases deep-links to it. - features/chat/: per-thread ragSource/enableRerank/ragTopK state in chat-runtime-store; Retrieval section in chat-settings-sheet with KB DropdownMenu (active highlight + per-row trash), thread doc list with Clear-thread-index button, RAG Top K slider, reranker toggle; chat-adapter retrieves before /v1/chat/completions and injects hits as a system block; shared-composer + button routes documents into pendingDocs (auto-uploads, send blocked while indexing).
113 lines
3.1 KiB
Python
113 lines
3.1 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Optional cross-encoder reranking stage.
|
|
|
|
Off by default. Callers opt in per request via ``enable_rerank`` on
|
|
``SearchRequest``. The reranker model is lazy-loaded on first opt-in
|
|
query and competes with the active chat model for GPU memory — keeping
|
|
it opt-in protects chat latency on smaller GPUs.
|
|
|
|
``sentence_transformers.CrossEncoder`` has no unsloth wrapper today;
|
|
this module loads it directly. A future ``FastCrossEncoder`` addition
|
|
to ``unsloth/models/sentence_transformer.py`` would slot in here by
|
|
replacing the import in ``_load``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import gc
|
|
import logging
|
|
import threading
|
|
from typing import Any
|
|
|
|
from utils.rag.config import RAG_RERANK_BATCH_SIZE, RAG_RERANKER_MODEL
|
|
|
|
from .retrieval import Hit
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_lock = threading.Lock()
|
|
_model: Any | None = None
|
|
_model_name: str | None = None
|
|
|
|
|
|
def _load(model_name: str) -> Any:
|
|
from sentence_transformers import CrossEncoder
|
|
|
|
logger.info("Loading RAG reranker: %s", model_name)
|
|
return CrossEncoder(model_name)
|
|
|
|
|
|
def get_reranker(model_name: str | None = None) -> Any:
|
|
global _model, _model_name
|
|
target = model_name or RAG_RERANKER_MODEL
|
|
with _lock:
|
|
if _model is None or _model_name != target:
|
|
unload()
|
|
_model = _load(target)
|
|
_model_name = target
|
|
return _model
|
|
|
|
|
|
def unload() -> None:
|
|
"""Drop the reranker reference and trigger a GC pass.
|
|
|
|
Useful when memory pressure is high — callers can free the
|
|
reranker without restarting the studio process. Next ``rerank``
|
|
call lazy-loads it again.
|
|
"""
|
|
global _model, _model_name
|
|
with _lock:
|
|
if _model is not None:
|
|
_model = None
|
|
_model_name = None
|
|
gc.collect()
|
|
try:
|
|
import torch
|
|
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
except ImportError:
|
|
pass
|
|
|
|
|
|
def rerank(
|
|
query: str,
|
|
pairs: list[tuple[Hit, str]],
|
|
*,
|
|
model_name: str | None = None,
|
|
top_k: int | None = None,
|
|
) -> list[Hit]:
|
|
"""Re-order ``pairs`` by CrossEncoder relevance to ``query``.
|
|
|
|
Each pair is ``(Hit, chunk_text)``. Returns Hits with the new
|
|
cross-encoder scores attached. If ``top_k`` is given, truncates.
|
|
"""
|
|
if not pairs:
|
|
return []
|
|
model = get_reranker(model_name)
|
|
inputs = [(query, text) for _, text in pairs]
|
|
scores = model.predict(
|
|
inputs,
|
|
batch_size = RAG_RERANK_BATCH_SIZE,
|
|
show_progress_bar = False,
|
|
)
|
|
ranked = sorted(
|
|
zip(pairs, scores),
|
|
key = lambda item: float(item[1]),
|
|
reverse = True,
|
|
)
|
|
out: list[Hit] = []
|
|
for (hit, _text), score in ranked:
|
|
out.append(
|
|
Hit(
|
|
chunk_id = hit.chunk_id,
|
|
score = float(score),
|
|
document_id = hit.document_id,
|
|
chunk_index = hit.chunk_index,
|
|
)
|
|
)
|
|
if top_k is not None:
|
|
out = out[:top_k]
|
|
return out
|