scriptoriumadmin/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2026-09-17 23:52:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

feat: add token_transformers text splitter using HuggingFace tokenizer

Browse source
This commit is contained in:
kela4 2026-04-25 20:08:41 +02:00
parent 4e2240aada
commit 25a0475598
6 changed files with 309 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
backend/open_webui/config.py
View file

@ -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(

15
backend/open_webui/main.py
View file

@ -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

137
backend/open_webui/routers/retrieval.py
View file

@ -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'))

32
src/lib/apis/retrieval/index.ts
View file

@ -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;

121
src/lib/components/admin/Settings/Documents.svelte
View file

@ -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();
});
</script>
@ -797,10 +827,69 @@
>
<option value="">{$i18n.t('Default')} ({$i18n.t('Character')})</option>
<option value="token">{$i18n.t('Token')} ({$i18n.t('Tiktoken')})</option>
<option value="token_transformers">{$i18n.t('Token')} ({$i18n.t('Transformers')})</option>
</select>
</div>
</div>
{#if RAGConfig.TEXT_SPLITTER === 'token_transformers'}
<div class="mb-2.5 flex w-full justify-between">
<div class="flex gap-1.5 w-full">
<div class="w-full">
<div class="self-center text-xs font-medium min-w-fit mb-1">
<Tooltip
placement="top-start"
content={$i18n.t(
'HuggingFace repository name of a model to load its tokenizer locally for exact token length calculation (e.g. sentence-transformers/all-MiniLM-L6-v2). Takes priority over the local embedding model\'s tokenizer when set. Required when using an external embedding API.'
)}
>
{$i18n.t('Tokenizer Model (HuggingFace Repo)')}
{#if RAG_EMBEDDING_ENGINE !== ''}
<span class="text-red-500 ml-0.5">*</span>
{/if}
</Tooltip>
</div>
<div class="self-center flex gap-2">
<input
class="w-full rounded-lg py-1.5 px-4 text-sm bg-gray-50 dark:text-gray-300 dark:bg-gray-850 outline-hidden"
type="text"
placeholder={$i18n.t('e.g. sentence-transformers/all-MiniLM-L6-v2')}
bind:value={RAGConfig.RAG_TOKENIZER_MODEL}
autocomplete="off"
/>
<button
class="px-2.5 bg-transparent text-gray-800 dark:bg-transparent dark:text-gray-100 rounded-lg transition"
on:click={tokenizerModelUpdateHandler}
disabled={updateTokenizerModelLoading || !RAGConfig.RAG_TOKENIZER_MODEL}
title={$i18n.t('Download tokenizer')}
>
{#if updateTokenizerModelLoading}
<div class="self-center">
<Spinner />
</div>
{:else}
<svg
xmlns="http://www.w3.org/2000/svg"
viewBox="0 0 16 16"
fill="currentColor"
class="w-4 h-4"
>
<path
d="M8.75 2.75a.75.75 0 0 0-1.5 0v5.69L5.03 6.22a.75.75 0 0 0-1.06 1.06l3.5 3.5a.75.75 0 0 0 1.06 0l3.5-3.5a.75.75 0 0 0-1.06-1.06L8.75 8.44V2.75Z"
/>
<path
d="M3.5 9.75a.75.75 0 0 0-1.5 0v1.5A2.75 2.75 0 0 0 4.75 14h6.5A2.75 2.75 0 0 0 14 11.25v-1.5a.75.75 0 0 0-1.5 0v1.5c0 .69-.56 1.25-1.25 1.25h-6.5c-.69 0-1.25-.56-1.25-1.25v-1.5Z"
/>
</svg>
{/if}
</button>
</div>
</div>
</div>
</div>
{/if}
<div class=" mb-2.5 flex w-full justify-between">
<div class=" self-center text-xs font-medium">
<Tooltip

5
src/lib/i18n/locales/en-US/translation.json
View file

@ -617,6 +617,7 @@
"don't install random tools from sources you don't trust.": "",
"Don't like the style": "",
"Done": "",
"Download tokenizer": "",
"Download": "",
"Download & Delete": "",
"Download as JSON": "",
@ -644,6 +645,7 @@
"e.g. my_filter": "",
"e.g. my_tools": "",
"e.g. pdf, docx, txt": "",
"e.g. sentence-transformers/all-MiniLM-L6-v2": "",
"e.g. Step-by-step instructions for code reviews": "",
"e.g. Tell me a fun fact": "",
"e.g. Tell me a fun fact about the Roman Empire": "",
@ -1072,6 +1074,7 @@
"Host": "",
"Hourly": "",
"Hourly Messages": "",
"HuggingFace repository name of a model to load its tokenizer locally for exact token length calculation (e.g. sentence-transformers/all-MiniLM-L6-v2). Takes priority over the local embedding model's tokenizer when set. Required when using an external embedding API.": "",
"How can I help you today?": "",
"How would you rate this response?": "",
"HTML": "",
@ -2070,6 +2073,8 @@
"Toggle Sidebar": "",
"Toggle status history": "",
"Toggle whether current connection is active.": "",
"Tokenizer Model (HuggingFace Repo)": "",
"Tokenizer Model required when using an external embedding engine.": "",
"Token": "",
"Token counts are estimates and may not reflect actual API usage": "",
"tokens": "",