From 25a0475598aacfd95d929b82fc15ed4573ad266c Mon Sep 17 00:00:00 2001 From: kela4 Date: Sat, 25 Apr 2026 20:08:41 +0200 Subject: [PATCH] feat: add token_transformers text splitter using HuggingFace tokenizer --- backend/open_webui/config.py | 15 ++ backend/open_webui/main.py | 15 ++ backend/open_webui/routers/retrieval.py | 137 ++++++++++++++++++ src/lib/apis/retrieval/index.ts | 32 ++++ .../admin/Settings/Documents.svelte | 121 ++++++++++++++-- src/lib/i18n/locales/en-US/translation.json | 5 + 6 files changed, 309 insertions(+), 16 deletions(-) diff --git a/backend/open_webui/config.py b/backend/open_webui/config.py index 06178d385c..9eedbb8f4f 100644 --- a/backend/open_webui/config.py +++ b/backend/open_webui/config.py @@ -3038,6 +3038,21 @@ TIKTOKEN_ENCODING_NAME = PersistentConfig( ) +RAG_TOKENIZER_MODEL = PersistentConfig( + 'RAG_TOKENIZER_MODEL', + 'rag.tokenizer_model', + os.environ.get('RAG_TOKENIZER_MODEL', ''), +) + +RAG_TOKENIZER_MODEL_AUTO_UPDATE = ( + not OFFLINE_MODE and os.environ.get('RAG_TOKENIZER_MODEL_AUTO_UPDATE', 'True').lower() == 'true' +) + +RAG_TOKENIZER_MODEL_TRUST_REMOTE_CODE = ( + os.environ.get('RAG_TOKENIZER_MODEL_TRUST_REMOTE_CODE', 'False').lower() == 'true' +) + + CHUNK_SIZE = PersistentConfig('CHUNK_SIZE', 'rag.chunk_size', int(os.environ.get('CHUNK_SIZE', '1000'))) CHUNK_MIN_SIZE_TARGET = PersistentConfig( diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py index af570af0af..75d240d700 100644 --- a/backend/open_webui/main.py +++ b/backend/open_webui/main.py @@ -113,6 +113,7 @@ from open_webui.routers.retrieval import ( get_reranking_function, get_ef, get_rf, + get_rag_tokenizer, ) @@ -306,6 +307,7 @@ from open_webui.config import ( PADDLEOCR_VL_BASE_URL, PADDLEOCR_VL_TOKEN, RAG_TEXT_SPLITTER, + RAG_TOKENIZER_MODEL, ENABLE_MARKDOWN_HEADER_TEXT_SPLITTER, TIKTOKEN_ENCODING_NAME, PDF_EXTRACT_IMAGES, @@ -1034,6 +1036,7 @@ app.state.config.MINERU_API_TIMEOUT = MINERU_API_TIMEOUT app.state.config.MINERU_PARAMS = MINERU_PARAMS app.state.config.TEXT_SPLITTER = RAG_TEXT_SPLITTER +app.state.config.RAG_TOKENIZER_MODEL = RAG_TOKENIZER_MODEL app.state.config.ENABLE_MARKDOWN_HEADER_TEXT_SPLITTER = ENABLE_MARKDOWN_HEADER_TEXT_SPLITTER app.state.config.TIKTOKEN_ENCODING_NAME = TIKTOKEN_ENCODING_NAME @@ -1147,6 +1150,7 @@ app.state.EMBEDDING_FUNCTION = None app.state.RERANKING_FUNCTION = None app.state.ef = None app.state.rf = None +app.state.rag_tokenizer = None app.state.YOUTUBE_LOADER_TRANSLATION = None @@ -1163,6 +1167,17 @@ try: ) else: app.state.rf = None + if app.state.config.RAG_TOKENIZER_MODEL: + tokenizer_model_name = str(app.state.config.RAG_TOKENIZER_MODEL) + log.info(f'Loading RAG tokenizer model: {tokenizer_model_name}') + app.state.rag_tokenizer = get_rag_tokenizer(tokenizer_model_name) + if app.state.rag_tokenizer is None: + log.error( + f'RAG tokenizer model \'{tokenizer_model_name}\' could not be loaded at startup. ' + f'Uploads using token_transformers splitter will fail until the model is available.' + ) + else: + log.info(f'RAG tokenizer model \'{tokenizer_model_name}\' loaded successfully') except Exception as e: log.error(f'Error updating models: {e}') pass diff --git a/backend/open_webui/routers/retrieval.py b/backend/open_webui/routers/retrieval.py index 01ef1d8886..40b459bf7c 100644 --- a/backend/open_webui/routers/retrieval.py +++ b/backend/open_webui/routers/retrieval.py @@ -111,6 +111,8 @@ from open_webui.config import ( DEFAULT_LOCALE, RAG_EMBEDDING_CONTENT_PREFIX, RAG_EMBEDDING_QUERY_PREFIX, + RAG_TOKENIZER_MODEL_AUTO_UPDATE, + RAG_TOKENIZER_MODEL_TRUST_REMOTE_CODE, ) from open_webui.env import ( DEVICE_TYPE, @@ -121,6 +123,7 @@ from open_webui.env import ( SENTENCE_TRANSFORMERS_CROSS_ENCODER_BACKEND, SENTENCE_TRANSFORMERS_CROSS_ENCODER_MODEL_KWARGS, SENTENCE_TRANSFORMERS_CROSS_ENCODER_SIGMOID_ACTIVATION_FUNCTION, + OFFLINE_MODE, ) from open_webui.constants import ERROR_MESSAGES @@ -136,6 +139,46 @@ log = logging.getLogger(__name__) ########################################## +def get_rag_tokenizer( + tokenizer_model: str, + auto_update: bool = RAG_TOKENIZER_MODEL_AUTO_UPDATE, + force_update: bool = False, +): + if not tokenizer_model: + return None + try: + from transformers import AutoTokenizer + + local_files_only = not (force_update or auto_update) + if OFFLINE_MODE: + local_files_only = True + + cache_dir = os.getenv('SENTENCE_TRANSFORMERS_HOME') or os.getenv('HF_HUB_CACHE') + + # Default to repo_id; AutoTokenizer will only pull tokenizer files when downloading + model_path = tokenizer_model + if not force_update: + # If the full snapshot is already cached (e.g. same model used for local embeddings), + # resolve it via get_model_path so we can use the local directory directly. + try: + candidate = get_model_path(tokenizer_model, update_model=False) + if os.path.exists(candidate): + model_path = candidate + except Exception: + if OFFLINE_MODE: + raise + + return AutoTokenizer.from_pretrained( + model_path, + cache_dir=cache_dir, + trust_remote_code=RAG_TOKENIZER_MODEL_TRUST_REMOTE_CODE, + local_files_only=local_files_only, + ) + except Exception as e: + log.error(f'Error loading tokenizer {tokenizer_model}: {e}') + return None + + def get_ef( engine: str, embedding_model: str, @@ -439,6 +482,35 @@ async def update_embedding_config(request: Request, form_data: EmbeddingModelUpd ) +class TokenizerModelUpdateForm(BaseModel): + RAG_TOKENIZER_MODEL: str + + +@router.post('/tokenizer/update') +async def update_tokenizer_model(request: Request, form_data: TokenizerModelUpdateForm, user=Depends(get_admin_user)): + tokenizer_model = form_data.RAG_TOKENIZER_MODEL.strip() + if not tokenizer_model: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail='RAG_TOKENIZER_MODEL is not set', + ) + + log.info(f'Downloading/updating tokenizer: {tokenizer_model}') + try: + tokenizer = get_rag_tokenizer(tokenizer_model, force_update=True) + if tokenizer is None: + raise ValueError(f'Failed to load tokenizer: {tokenizer_model}') + request.app.state.rag_tokenizer = tokenizer + request.app.state.config.RAG_TOKENIZER_MODEL = tokenizer_model + return {'status': True, 'RAG_TOKENIZER_MODEL': tokenizer_model} + except Exception as e: + log.exception(f'Problem updating tokenizer model: {e}') + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=ERROR_MESSAGES.DEFAULT(e), + ) + + @router.get('/config') async def get_rag_config(request: Request, user=Depends(get_admin_user)): return { @@ -497,6 +569,7 @@ async def get_rag_config(request: Request, user=Depends(get_admin_user)): 'RAG_EXTERNAL_RERANKER_TIMEOUT': request.app.state.config.RAG_EXTERNAL_RERANKER_TIMEOUT, # Chunking settings 'TEXT_SPLITTER': request.app.state.config.TEXT_SPLITTER, + 'RAG_TOKENIZER_MODEL': request.app.state.config.RAG_TOKENIZER_MODEL, 'ENABLE_MARKDOWN_HEADER_TEXT_SPLITTER': request.app.state.config.ENABLE_MARKDOWN_HEADER_TEXT_SPLITTER, 'CHUNK_SIZE': request.app.state.config.CHUNK_SIZE, 'CHUNK_MIN_SIZE_TARGET': request.app.state.config.CHUNK_MIN_SIZE_TARGET, @@ -708,6 +781,7 @@ class ConfigForm(BaseModel): # Chunking settings TEXT_SPLITTER: Optional[str] = None + RAG_TOKENIZER_MODEL: Optional[str] = None ENABLE_MARKDOWN_HEADER_TEXT_SPLITTER: Optional[bool] = None CHUNK_SIZE: Optional[int] = None CHUNK_MIN_SIZE_TARGET: Optional[int] = None @@ -1023,6 +1097,9 @@ async def update_rag_config(request: Request, form_data: ConfigForm, user=Depend request.app.state.config.CHUNK_OVERLAP = ( form_data.CHUNK_OVERLAP if form_data.CHUNK_OVERLAP is not None else request.app.state.config.CHUNK_OVERLAP ) + if form_data.RAG_TOKENIZER_MODEL is not None and form_data.RAG_TOKENIZER_MODEL != str(request.app.state.config.RAG_TOKENIZER_MODEL): + request.app.state.config.RAG_TOKENIZER_MODEL = form_data.RAG_TOKENIZER_MODEL + request.app.state.rag_tokenizer = None # invalidate cache so it reloads on next ingestion # File upload settings # Empty string means "clear to None" (unlimited/no compression), @@ -1283,6 +1360,28 @@ def can_merge_chunks(a: Document, b: Document) -> bool: return True +def get_transformers_tokenizer_for_text_splitter(request: Request): + tokenizer_model = str(request.app.state.config.RAG_TOKENIZER_MODEL).strip() + if tokenizer_model: + if request.app.state.rag_tokenizer is not None: + return request.app.state.rag_tokenizer + + tokenizer = get_rag_tokenizer(tokenizer_model) + if tokenizer is None: + raise ValueError( + f"RAG_TOKENIZER_MODEL is set to '{tokenizer_model}' but failed to load — " + f"check the application logs for details. " + f"Ensure the model name is correct and the model can be downloaded or is already cached locally." + ) + request.app.state.rag_tokenizer = tokenizer + return tokenizer + + if request.app.state.ef is not None and hasattr(request.app.state.ef, "tokenizer"): + return request.app.state.ef.tokenizer + + return None + + def merge_docs_to_target_size( request: Request, chunks: list[Document], @@ -1304,6 +1403,12 @@ def merge_docs_to_target_size( if request.app.state.config.TEXT_SPLITTER == 'token': encoding = tiktoken.get_encoding(str(request.app.state.config.TIKTOKEN_ENCODING_NAME)) measure_chunk_size = lambda text: len(encoding.encode(text)) + elif request.app.state.config.TEXT_SPLITTER == 'token_transformers': + tokenizer = get_transformers_tokenizer_for_text_splitter(request) + if tokenizer is not None: + measure_chunk_size = lambda text: len(tokenizer.encode(text, add_special_tokens=False)) + else: + log.warning('token_transformers tokenizer unavailable in merge_docs_to_target_size, falling back to character counting') processed_chunks: list[Document] = [] @@ -1446,6 +1551,38 @@ def save_docs_to_vector_db( add_start_index=True, ) docs = text_splitter.split_documents(docs) + elif request.app.state.config.TEXT_SPLITTER == 'token_transformers': + tokenizer = get_transformers_tokenizer_for_text_splitter(request) + if tokenizer is None: + raise ValueError( + 'token_transformers splitter requires a local embedding model or RAG_TOKENIZER_MODEL to be set' + ) + + model_max_length = getattr(tokenizer, 'model_max_length', None) + chunk_size = request.app.state.config.CHUNK_SIZE + if model_max_length and model_max_length <= 100_000: + num_special = tokenizer.num_special_tokens_to_add(pair=False) + effective_max = model_max_length - num_special + if chunk_size > effective_max: + log.warning( + f'token_transformers splitter: CHUNK_SIZE={chunk_size} exceeds ' + f'effective limit of {effective_max} ' + f'(model_max_length={model_max_length} minus {num_special} special tokens) — ' + f'chunks will be silently truncated during embedding (local embedding model) ' + f'or embedding will fail (remote API embedding model). ' + f'Consider reducing CHUNK_SIZE to at most {effective_max}.' + ) + + def token_length(text): + return len(tokenizer.encode(text, add_special_tokens=False)) + + text_splitter = RecursiveCharacterTextSplitter( + chunk_size=chunk_size, + chunk_overlap=request.app.state.config.CHUNK_OVERLAP, + length_function=token_length, + add_start_index=True, + ) + docs = text_splitter.split_documents(docs) else: raise ValueError(ERROR_MESSAGES.DEFAULT('Invalid text splitter')) diff --git a/src/lib/apis/retrieval/index.ts b/src/lib/apis/retrieval/index.ts index a84e7b6822..f621d7ba13 100644 --- a/src/lib/apis/retrieval/index.ts +++ b/src/lib/apis/retrieval/index.ts @@ -90,6 +90,38 @@ export const updateRAGConfig = async (token: string, payload: RAGConfigForm) => return res; }; +type TokenizerModelUpdateForm = { + RAG_TOKENIZER_MODEL: string; +}; + +export const updateTokenizerModel = async (token: string, payload: TokenizerModelUpdateForm) => { + let error = null; + + const res = await fetch(`${RETRIEVAL_API_BASE_URL}/tokenizer/update`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: `Bearer ${token}` + }, + body: JSON.stringify({ ...payload }) + }) + .then(async (res) => { + if (!res.ok) throw await res.json(); + return res.json(); + }) + .catch((err) => { + console.error(err); + error = err.detail; + return null; + }); + + if (error) { + throw error; + } + + return res; +}; + export const getQuerySettings = async (token: string) => { let error = null; diff --git a/src/lib/components/admin/Settings/Documents.svelte b/src/lib/components/admin/Settings/Documents.svelte index e173c74969..93372526cd 100644 --- a/src/lib/components/admin/Settings/Documents.svelte +++ b/src/lib/components/admin/Settings/Documents.svelte @@ -14,7 +14,8 @@ getRerankingConfig, updateRerankingConfig, getRAGConfig, - updateRAGConfig + updateRAGConfig, + updateTokenizerModel } from '$lib/apis/retrieval'; import { reindexKnowledgeFiles } from '$lib/apis/knowledge'; @@ -33,6 +34,7 @@ let updateEmbeddingModelLoading = false; let updateRerankingModelLoading = false; + let updateTokenizerModelLoading = false; let showResetConfirm = false; let showResetUploadDirConfirm = false; @@ -141,6 +143,22 @@ } }; + const tokenizerModelUpdateHandler = async () => { + updateTokenizerModelLoading = true; + const res = await updateTokenizerModel(localStorage.token, { + RAG_TOKENIZER_MODEL: RAGConfig.RAG_TOKENIZER_MODEL + }).catch(async (error) => { + toast.error(`${error}`); + await setRAGConfig(); + return null; + }); + updateTokenizerModelLoading = false; + + if (res) { + toast.success($i18n.t('Success')); + } + }; + const submitHandler = async () => { if ( RAGConfig.CONTENT_EXTRACTION_ENGINE === 'external' && @@ -201,6 +219,15 @@ return; } + if ( + RAGConfig.TEXT_SPLITTER === 'token_transformers' && + RAG_EMBEDDING_ENGINE !== '' && + !RAGConfig.RAG_TOKENIZER_MODEL?.trim() + ) { + toast.error($i18n.t('Tokenizer Model required when using an external embedding engine.')); + return; + } + if (!RAGConfig.BYPASS_EMBEDDING_AND_RETRIEVAL) { await embeddingModelUpdateHandler(); } @@ -249,6 +276,23 @@ dispatch('save'); }; + const setRAGConfig = async () => { + const config = await getRAGConfig(localStorage.token); + config.ALLOWED_FILE_EXTENSIONS = (config?.ALLOWED_FILE_EXTENSIONS ?? []).join(', '); + + config.DOCLING_PARAMS = + typeof config.DOCLING_PARAMS === 'object' + ? JSON.stringify(config.DOCLING_PARAMS ?? {}, null, 2) + : config.DOCLING_PARAMS; + + config.MINERU_PARAMS = + typeof config.MINERU_PARAMS === 'object' + ? JSON.stringify(config.MINERU_PARAMS ?? {}, null, 2) + : config.MINERU_PARAMS; + + RAGConfig = config; + }; + const setEmbeddingConfig = async () => { const embeddingConfig = await getEmbeddingConfig(localStorage.token); @@ -272,21 +316,7 @@ }; onMount(async () => { await setEmbeddingConfig(); - - const config = await getRAGConfig(localStorage.token); - config.ALLOWED_FILE_EXTENSIONS = (config?.ALLOWED_FILE_EXTENSIONS ?? []).join(', '); - - config.DOCLING_PARAMS = - typeof config.DOCLING_PARAMS === 'object' - ? JSON.stringify(config.DOCLING_PARAMS ?? {}, null, 2) - : config.DOCLING_PARAMS; - - config.MINERU_PARAMS = - typeof config.MINERU_PARAMS === 'object' - ? JSON.stringify(config.MINERU_PARAMS ?? {}, null, 2) - : config.MINERU_PARAMS; - - RAGConfig = config; + await setRAGConfig(); }); @@ -797,10 +827,69 @@ > + + {#if RAGConfig.TEXT_SPLITTER === 'token_transformers'} +
+
+
+
+ + {$i18n.t('Tokenizer Model (HuggingFace Repo)')} + {#if RAG_EMBEDDING_ENGINE !== ''} + * + {/if} + +
+
+ + + +
+
+
+
+ {/if} +