This commit is contained in:
Timothy Jaeryang Baek
2026-06-29 12:26:23 -05:00
parent 248315de14
commit bb6b2db88b
3 changed files with 81 additions and 4 deletions
+3
View File
@@ -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,
+57 -4
View File
@@ -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