90 lines
3.2 KiB
Python
90 lines
3.2 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
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
|
|
|
|
def _env_int(name: str, default: int) -> int:
|
|
raw = os.environ.get(name, "").strip()
|
|
if not raw:
|
|
return default
|
|
try:
|
|
return int(raw)
|
|
except ValueError:
|
|
return default
|
|
|
|
|
|
def _env_float(name: str, default: float) -> float:
|
|
raw = os.environ.get(name, "").strip()
|
|
if not raw:
|
|
return default
|
|
try:
|
|
return float(raw)
|
|
except ValueError:
|
|
return default
|
|
|
|
|
|
RAG_EMBEDDING_MODEL: str = (
|
|
os.environ.get("UNSLOTH_RAG_EMBEDDING_MODEL", "").strip()
|
|
or "BAAI/bge-small-en-v1.5"
|
|
)
|
|
|
|
# Phase 3: default embedders per (mode, chunking_strategy). Ingestion in
|
|
# Phase 3B-late and Phase 3B-multimodal looks the embedder up here at job
|
|
# start, falling back to RAG_EMBEDDING_MODEL (above) for legacy KBs that
|
|
# pre-date the columns. The (multimodal, late) combo is intentionally
|
|
# absent — no public open-weight embedder supports both at once, and
|
|
# routes/rag.py rejects the combo with a 400 at KB create time.
|
|
RAG_EMBEDDER_MATRIX: dict[tuple[str, str], str] = {
|
|
("text", "standard"): "BAAI/bge-small-en-v1.5",
|
|
("text", "late"): "nomic-ai/nomic-embed-text-v1.5",
|
|
# Qwen3-VL-Embedding-2B: 2B-param multimodal embedder built on
|
|
# Qwen3 (not CLIP), so no 77-token text cap — long chunks embed
|
|
# losslessly. Loads through vanilla SentenceTransformer with
|
|
# trust_remote_code, no shim. 2048-d shared text/image space.
|
|
# Trade-off: ~4 GB download / VRAM vs BGE-VL-base's ~600 MB.
|
|
("multimodal", "standard"): "Qwen/Qwen3-VL-Embedding-2B",
|
|
}
|
|
|
|
|
|
def resolve_embedder(mode: str, chunking_strategy: str) -> str:
|
|
"""Look up the default embedder for a (mode, chunking_strategy) pair.
|
|
|
|
Unknown combos fall back to the legacy single default so old KBs
|
|
keep working. Callers that explicitly require the new matrix
|
|
behaviour (Phase 3B paths) should validate the inputs before
|
|
calling.
|
|
"""
|
|
return RAG_EMBEDDER_MATRIX.get(
|
|
(mode, chunking_strategy),
|
|
RAG_EMBEDDING_MODEL,
|
|
)
|
|
|
|
|
|
RAG_CHUNK_SIZE: int = _env_int("UNSLOTH_RAG_CHUNK_SIZE", 512)
|
|
RAG_CHUNK_OVERLAP: int = _env_int("UNSLOTH_RAG_CHUNK_OVERLAP", 64)
|
|
|
|
RAG_TOP_K_BM25: int = _env_int("UNSLOTH_RAG_TOP_K_BM25", 30)
|
|
RAG_TOP_K_DENSE: int = _env_int("UNSLOTH_RAG_TOP_K_DENSE", 30)
|
|
RAG_TOP_K_HYBRID: int = _env_int("UNSLOTH_RAG_TOP_K_HYBRID", 10)
|
|
|
|
RAG_RRF_K: int = _env_int("UNSLOTH_RAG_RRF_K", 60)
|
|
|
|
RAG_MAX_UPLOAD_MB: int = _env_int("UNSLOTH_RAG_MAX_UPLOAD_MB", 50)
|
|
|
|
RAG_EMBED_BATCH_SIZE: int = _env_int("UNSLOTH_RAG_EMBED_BATCH_SIZE", 32)
|
|
|
|
# Reranking is off by default. The CrossEncoder runs on GPU and competes
|
|
# with the active chat model — callers opt in per-request via
|
|
# `enable_rerank` on SearchRequest.
|
|
RAG_RERANKER_MODEL: str = (
|
|
os.environ.get("UNSLOTH_RAG_RERANKER_MODEL", "").strip() or "BAAI/bge-reranker-base"
|
|
)
|
|
RAG_RERANK_CANDIDATE_K: int = _env_int("UNSLOTH_RAG_RERANK_CANDIDATE_K", 50)
|
|
RAG_RERANK_BATCH_SIZE: int = _env_int("UNSLOTH_RAG_RERANK_BATCH_SIZE", 16)
|
|
|
|
RAG_UPLOAD_EXTS: frozenset[str] = frozenset(
|
|
{".pdf", ".txt", ".md", ".markdown", ".docx", ".html", ".htm"}
|
|
)
|