mirror of
https://github.com/open-webui/open-webui.git
synced 2026-07-20 20:52:23 -05:00
refac
This commit is contained in:
@@ -936,6 +936,8 @@ PDF_LOADER_MODE = os.getenv('PDF_LOADER_MODE', 'page')
|
||||
RAG_EMBEDDING_MODEL = os.getenv('RAG_EMBEDDING_MODEL', 'sentence-transformers/all-MiniLM-L6-v2')
|
||||
log.info(f'Embedding model set: {RAG_EMBEDDING_MODEL}')
|
||||
|
||||
RAG_TOKENIZER_MODEL = os.getenv('RAG_TOKENIZER_MODEL', '')
|
||||
|
||||
RAG_EMBEDDING_MODEL_AUTO_UPDATE = (
|
||||
not OFFLINE_MODE and os.getenv('RAG_EMBEDDING_MODEL_AUTO_UPDATE', 'True').lower() == 'true'
|
||||
)
|
||||
@@ -2751,6 +2753,7 @@ DEFAULT_CONFIG = {
|
||||
'rag.pdf_extract_images': PDF_EXTRACT_IMAGES,
|
||||
'rag.pdf_loader_mode': PDF_LOADER_MODE,
|
||||
'rag.embedding_model': RAG_EMBEDDING_MODEL,
|
||||
'rag.tokenizer_model': RAG_TOKENIZER_MODEL,
|
||||
'rag.embedding_batch_size': RAG_EMBEDDING_BATCH_SIZE,
|
||||
'rag.enable_async_embedding': ENABLE_ASYNC_EMBEDDING,
|
||||
'rag.embedding_concurrent_requests': RAG_EMBEDDING_CONCURRENT_REQUESTS,
|
||||
|
||||
@@ -346,6 +346,7 @@ RETRIEVAL_CONFIG_KEYS = {
|
||||
'RAG_EMBEDDING_CONCURRENT_REQUESTS': 'rag.embedding_concurrent_requests',
|
||||
'RAG_EMBEDDING_ENGINE': 'rag.embedding_engine',
|
||||
'RAG_EMBEDDING_MODEL': 'rag.embedding_model',
|
||||
'RAG_TOKENIZER_MODEL': 'rag.tokenizer_model',
|
||||
'RAG_EXTERNAL_RERANKER_API_KEY': 'rag.external_reranker_api_key',
|
||||
'RAG_EXTERNAL_RERANKER_TIMEOUT': 'rag.external_reranker_timeout',
|
||||
'RAG_EXTERNAL_RERANKER_URL': 'rag.external_reranker_url',
|
||||
@@ -673,6 +674,7 @@ async def get_rag_config(request: Request, user=Depends(get_admin_user)):
|
||||
'RAG_EXTERNAL_RERANKER_TIMEOUT': config.RAG_EXTERNAL_RERANKER_TIMEOUT,
|
||||
# Chunking settings
|
||||
'TEXT_SPLITTER': config.TEXT_SPLITTER,
|
||||
'RAG_TOKENIZER_MODEL': config.RAG_TOKENIZER_MODEL,
|
||||
'ENABLE_MARKDOWN_HEADER_TEXT_SPLITTER': config.ENABLE_MARKDOWN_HEADER_TEXT_SPLITTER,
|
||||
'CHUNK_SIZE': config.CHUNK_SIZE,
|
||||
'CHUNK_MIN_SIZE_TARGET': config.CHUNK_MIN_SIZE_TARGET,
|
||||
@@ -906,6 +908,7 @@ class ConfigForm(BaseModel):
|
||||
|
||||
# Chunking settings
|
||||
TEXT_SPLITTER: str | None = None
|
||||
RAG_TOKENIZER_MODEL: str | None = None
|
||||
ENABLE_MARKDOWN_HEADER_TEXT_SPLITTER: bool | None = None
|
||||
CHUNK_SIZE: int | None = None
|
||||
CHUNK_MIN_SIZE_TARGET: int | None = None
|
||||
@@ -1232,6 +1235,11 @@ async def update_rag_config(request: Request, form_data: ConfigForm, user=Depend
|
||||
config.CHUNK_OVERLAP = (
|
||||
form_data.CHUNK_OVERLAP if form_data.CHUNK_OVERLAP is not None else config.CHUNK_OVERLAP
|
||||
)
|
||||
config.RAG_TOKENIZER_MODEL = (
|
||||
form_data.RAG_TOKENIZER_MODEL.strip()
|
||||
if form_data.RAG_TOKENIZER_MODEL is not None
|
||||
else config.RAG_TOKENIZER_MODEL
|
||||
)
|
||||
|
||||
# File upload settings
|
||||
# Empty string means "clear to None" (unlimited/no compression),
|
||||
@@ -1405,6 +1413,7 @@ async def update_rag_config(request: Request, form_data: ConfigForm, user=Depend
|
||||
'RAG_EXTERNAL_RERANKER_TIMEOUT': config.RAG_EXTERNAL_RERANKER_TIMEOUT,
|
||||
# Chunking settings
|
||||
'TEXT_SPLITTER': config.TEXT_SPLITTER,
|
||||
'RAG_TOKENIZER_MODEL': config.RAG_TOKENIZER_MODEL,
|
||||
'CHUNK_SIZE': config.CHUNK_SIZE,
|
||||
'CHUNK_MIN_SIZE_TARGET': config.CHUNK_MIN_SIZE_TARGET,
|
||||
'ENABLE_MARKDOWN_HEADER_TEXT_SPLITTER': config.ENABLE_MARKDOWN_HEADER_TEXT_SPLITTER,
|
||||
@@ -1538,10 +1547,7 @@ def merge_docs_to_target_size(
|
||||
if min_size <= 0:
|
||||
return chunks
|
||||
|
||||
measure: Callable[[str], int] = len
|
||||
if config.TEXT_SPLITTER == 'token':
|
||||
encoding = tiktoken.get_encoding(str(config.TIKTOKEN_ENCODING_NAME))
|
||||
measure = lambda text: len(encoding.encode(text, disallowed_special=()))
|
||||
measure = get_splitter_length_function(request, config)
|
||||
|
||||
def _merge_backward(result: list[Document], content: str, chunk: Document) -> bool:
|
||||
"""Try to append content into the last emitted chunk. Returns True on success."""
|
||||
@@ -1594,6 +1600,43 @@ def merge_docs_to_target_size(
|
||||
return result
|
||||
|
||||
|
||||
def get_transformers_tokenizer(request: Request, config: RetrievalConfig):
|
||||
if config.RAG_TOKENIZER_MODEL:
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
tokenizer_model = config.RAG_TOKENIZER_MODEL
|
||||
if not os.path.exists(tokenizer_model) and '/' not in tokenizer_model:
|
||||
tokenizer_model = f'sentence-transformers/{tokenizer_model}'
|
||||
|
||||
return AutoTokenizer.from_pretrained(
|
||||
tokenizer_model,
|
||||
cache_dir=os.getenv('SENTENCE_TRANSFORMERS_HOME') or os.getenv('HF_HUB_CACHE'),
|
||||
trust_remote_code=RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE,
|
||||
local_files_only=not RAG_EMBEDDING_MODEL_AUTO_UPDATE,
|
||||
)
|
||||
|
||||
tokenizer = getattr(getattr(request.app.state, 'ef', None), 'tokenizer', None)
|
||||
if tokenizer is not None:
|
||||
return tokenizer
|
||||
|
||||
raise ValueError('Tokenizer model required for Token (Transformers) text splitter')
|
||||
|
||||
|
||||
def get_splitter_length_function(
|
||||
request: Request,
|
||||
config: RetrievalConfig,
|
||||
) -> Callable[[str], int]:
|
||||
if config.TEXT_SPLITTER == 'token':
|
||||
encoding = tiktoken.get_encoding(str(config.TIKTOKEN_ENCODING_NAME))
|
||||
return lambda text: len(encoding.encode(text, disallowed_special=()))
|
||||
|
||||
if config.TEXT_SPLITTER == 'token_transformers':
|
||||
tokenizer = get_transformers_tokenizer(request, config)
|
||||
return lambda text: len(tokenizer.encode(text))
|
||||
|
||||
return len
|
||||
|
||||
|
||||
def save_docs_to_vector_db(
|
||||
request: Request,
|
||||
docs,
|
||||
@@ -1694,6 +1737,16 @@ def save_docs_to_vector_db(
|
||||
add_start_index=True,
|
||||
)
|
||||
docs = text_splitter.split_documents(docs)
|
||||
elif config.TEXT_SPLITTER == 'token_transformers':
|
||||
log.info('Using transformers token text splitter')
|
||||
|
||||
text_splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=config.CHUNK_SIZE,
|
||||
chunk_overlap=config.CHUNK_OVERLAP,
|
||||
length_function=get_splitter_length_function(request, config),
|
||||
add_start_index=True,
|
||||
)
|
||||
docs = text_splitter.split_documents(docs)
|
||||
else:
|
||||
raise ValueError(ERROR_MESSAGES.DEFAULT('Invalid text splitter'))
|
||||
|
||||
|
||||
@@ -318,6 +318,7 @@
|
||||
: config.EXTERNAL_DOCUMENT_LOADER_HEADERS;
|
||||
|
||||
config.MINERU_FILE_EXTENSIONS = (config?.MINERU_FILE_EXTENSIONS ?? ['pdf']).join(', ');
|
||||
config.RAG_TOKENIZER_MODEL = config?.RAG_TOKENIZER_MODEL ?? '';
|
||||
|
||||
RAGConfig = config;
|
||||
});
|
||||
@@ -910,10 +911,30 @@
|
||||
>
|
||||
<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 flex-col w-full justify-between">
|
||||
<div class="self-center text-xs font-medium min-w-fit mb-1 w-full">
|
||||
{$i18n.t('Tokenizer Model')}
|
||||
</div>
|
||||
<div class="self-center w-full">
|
||||
<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"
|
||||
placeholder={$i18n.t('Enter Tokenizer Model')}
|
||||
bind:value={RAGConfig.RAG_TOKENIZER_MODEL}
|
||||
autocomplete="off"
|
||||
required={RAG_EMBEDDING_ENGINE !== ''}
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
{/if}
|
||||
|
||||
<div class=" mb-2.5 flex w-full justify-between">
|
||||
<div class=" self-center text-xs font-medium">
|
||||
<Tooltip
|
||||
|
||||
Reference in New Issue
Block a user