diff --git a/backend/apps/rag/main.py b/backend/apps/rag/main.py index 15de0cc45..7bb9ee0ee 100644 --- a/backend/apps/rag/main.py +++ b/backend/apps/rag/main.py @@ -44,7 +44,7 @@ from apps.web.models.documents import ( DocumentResponse, ) -from apps.rag.utils import query_doc, query_collection, embedding_model_get_path +from apps.rag.utils import query_doc, query_collection, get_embedding_model_path from utils.misc import ( calculate_sha256, @@ -77,10 +77,14 @@ app.state.PDF_EXTRACT_IMAGES = False app.state.CHUNK_SIZE = CHUNK_SIZE app.state.CHUNK_OVERLAP = CHUNK_OVERLAP app.state.RAG_TEMPLATE = RAG_TEMPLATE + + app.state.RAG_EMBEDDING_MODEL = RAG_EMBEDDING_MODEL -app.state.RAG_EMBEDDING_MODEL_PATH = embedding_model_get_path( +app.state.RAG_EMBEDDING_MODEL_PATH = get_embedding_model_path( app.state.RAG_EMBEDDING_MODEL, RAG_EMBEDDING_MODEL_AUTO_UPDATE ) + + app.state.TOP_K = 4 app.state.sentence_transformer_ef = ( @@ -148,7 +152,7 @@ async def update_embedding_model( ) try: - app.state.RAG_EMBEDDING_MODEL_PATH = embedding_model_get_path( + app.state.RAG_EMBEDDING_MODEL_PATH = get_embedding_model_path( app.state.RAG_EMBEDDING_MODEL, True ) app.state.sentence_transformer_ef = ( diff --git a/backend/apps/rag/utils.py b/backend/apps/rag/utils.py index 866759d9c..7bbfe0b88 100644 --- a/backend/apps/rag/utils.py +++ b/backend/apps/rag/utils.py @@ -192,21 +192,21 @@ def rag_messages(docs, messages, template, k, embedding_function): return messages -def embedding_model_get_path( +def get_embedding_model_path( embedding_model: str, update_embedding_model: bool = False ): # Construct huggingface_hub kwargs with local_files_only to return the snapshot path cache_dir = os.getenv("SENTENCE_TRANSFORMERS_HOME") + local_files_only = not update_embedding_model + snapshot_kwargs = { "cache_dir": cache_dir, "local_files_only": local_files_only, } - log.debug(f"SENTENCE_TRANSFORMERS_HOME cache_dir: {cache_dir}") log.debug(f"embedding_model: {embedding_model}") - log.debug(f"update_embedding_model: {update_embedding_model}") - log.debug(f"local_files_only: {local_files_only}") + log.debug(f"snapshot_kwargs: {snapshot_kwargs}") # Inspiration from upstream sentence_transformers if ( diff --git a/src/lib/components/documents/Settings/General.svelte b/src/lib/components/documents/Settings/General.svelte index cacc593c5..17bd3ad04 100644 --- a/src/lib/components/documents/Settings/General.svelte +++ b/src/lib/components/documents/Settings/General.svelte @@ -21,8 +21,8 @@ export let saveHandler: Function; - let loading = false; - let loading1 = false; + let scanDirLoading = false; + let updateEmbeddingModelLoading = false; let showResetConfirm = false; @@ -35,14 +35,12 @@ k: 4 }; - let embeddingModel = { - embedding_model: '' - }; + let embeddingModel = ''; const scanHandler = async () => { - loading = true; + scanDirLoading = true; const res = await scanDocs(localStorage.token); - loading = false; + scanDirLoading = false; if (res) { await documents.set(await getDocs(localStorage.token)); @@ -51,7 +49,7 @@ }; const embeddingModelUpdateHandler = async () => { - if (embeddingModel.embedding_model.split('/').length - 1 > 1) { + if (embeddingModel.split('/').length - 1 > 1) { toast.error( $i18n.t( 'Model filesystem path detected. Model shortname is required for update, cannot continue.' @@ -60,11 +58,11 @@ return; } - console.log('Update embedding model attempt:', embeddingModel.embedding_model); + console.log('Update embedding model attempt:', embeddingModel); - loading1 = true; - const res = await updateEmbeddingModel(localStorage.token, embeddingModel); - loading1 = false; + updateEmbeddingModelLoading = true; + const res = await updateEmbeddingModel(localStorage.token, { embedding_model: embeddingModel }); + updateEmbeddingModelLoading = false; if (res) { console.log('embeddingModelUpdateHandler:', res); @@ -101,7 +99,9 @@ chunkOverlap = res.chunk.chunk_overlap; } - embeddingModel = await getEmbeddingModel(localStorage.token); + const embeddingModelConfig = await getEmbeddingModel(localStorage.token); + + embeddingModel = embeddingModelConfig.embedding_model; querySettings = await getQuerySettings(localStorage.token); }); @@ -124,7 +124,7 @@ - - - + -
+
+ {$i18n.t( + 'Warning: If you update or change your embedding model, you will need to re-import all documents.' + )} +
-
-
{$i18n.t('Chunk Params')}
+
-
-
-
{$i18n.t('Chunk Size')}
+
+
{$i18n.t('Chunk Params')}
-
- +
+
+
{$i18n.t('Chunk Size')}
+ +
+ +
+
+ +
+
+ {$i18n.t('Chunk Overlap')} +
+ +
+ +
+
+
+ +
+
+
{$i18n.t('PDF Extract Images (OCR)')}
+ + +
-
-
{$i18n.t('Chunk Overlap')}
+
+
{$i18n.t('Query Params')}
-
- -
-
-
+
+
+
{$i18n.t('Top K')}
-
-
-
{$i18n.t('PDF Extract Images (OCR)')}
+
+ +
+
- -
-
-
- -
-
{$i18n.t('Query Params')}
- -
-
-
{$i18n.t('Top K')}
- -
- -
-
- - +
+ +
+
{$i18n.t('RAG Template')}
+