refac

f4b87ecb · Timothy J. Baek · 48aad655 · f4b87ecb · f4b87ecb · f4b87ecb
Commit f4b87ecb authored Apr 10, 2024 by Timothy J. Baek
3 changed files
--- 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 = (

--- 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"snapshot_kwargs: {snapshot_kwargs}")
-    log.debug(f"local_files_only: {local_files_only}")
    # Inspiration from upstream sentence_transformers
    if (

--- 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 scanDirLoading = false;
-	let loading1 = false;
+	let updateEmbeddingModelLoading = false;
 	let showResetConfirm = false;
@@ -35,14 +35,12 @@
 		k: 4
 	};
-	let embeddingModel = {
+	let embeddingModel = '';
-		embedding_model: ''
-	};
 	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;
+		updateEmbeddingModelLoading = true;
-		const res = await updateEmbeddingModel(localStorage.token, embeddingModel);
+		const res = await updateEmbeddingModel(localStorage.token, { embedding_model: embeddingModel });
-		loading1 = false;
+		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 @@
 				</div>
 				<button
-					class=" self-center text-xs p-1 px-3 bg-gray-100 dark:bg-gray-800 dark:hover:bg-gray-700 rounded flex flex-row space-x-1 items-center {loading
+					class=" self-center text-xs p-1 px-3 bg-gray-100 dark:bg-gray-800 dark:hover:bg-gray-700 rounded flex flex-row space-x-1 items-center {scanDirLoading
 						? ' cursor-not-allowed'
 						: ''}"
 					on:click={() => {
@@ -132,24 +132,11 @@
 						console.log('check');
 					}}
 					type="button"
-					disabled={loading}
+					disabled={scanDirLoading}
 				>
 					<div class="self-center font-medium">{$i18n.t('Scan')}</div>
-					<!-- <svg
+					{#if scanDirLoading}
-						xmlns="http://www.w3.org/2000/svg"
-						viewBox="0 0 16 16"
-						fill="currentColor"
-						class="w-3 h-3"
-					>
-						<path
-							fill-rule="evenodd"
-							d="M13.836 2.477a.75.75 0 0 1 .75.75v3.182a.75.75 0 0 1-.75.75h-3.182a.75.75 0 0 1 0-1.5h1.37l-.84-.841a4.5 4.5 0 0 0-7.08.932.75.75 0 0 1-1.3-.75 6 6 0 0 1 9.44-1.242l.842.84V3.227a.75.75 0 0 1 .75-.75Zm-.911 7.5A.75.75 0 0 1 13.199 11a6 6 0 0 1-9.44 1.241l-.84-.84v1.371a.75.75 0 0 1-1.5 0V9.591a.75.75 0 0 1 .75-.75H5.35a.75.75 0 0 1 0 1.5H3.98l.841.841a4.5 4.5 0 0 0 7.08-.932.75.75 0 0 1 1.025-.273Z"
-							clip-rule="evenodd"
-						/>
-					</svg> -->
-					{#if loading}
 						<div class="ml-3 self-center">
 							<svg
 								class=" w-3 h-3"
@@ -182,36 +169,30 @@
 		<hr class=" dark:border-gray-700" />
+		<div class="space-y-2">
 			<div>
-			<div class="  flex w-full justify-between">
+				<div class=" mb-2 text-sm font-medium">{$i18n.t('Update Embedding Model')}</div>
-				<Tooltip content={$i18n.t('Embedding model: {{embedding_model}}', embeddingModel)}>
+				<div class="flex w-full">
-					<div class=" self-center text-xs font-medium">
+					<div class="flex-1 mr-2">
-						{$i18n.t('Update embedding model {{embedding_model}}', {
+						<input
-							embedding_model: embeddingModel.embedding_model.slice(-40)
+							class="w-full rounded-lg py-2 px-4 text-sm dark:text-gray-300 dark:bg-gray-850 outline-none"
+							placeholder={$i18n.t('Update embedding model (e.g. {{model}})', {
+								model: embeddingModel.slice(-40)
 							})}
+							bind:value={embeddingModel}
+						/>
 					</div>
-				</Tooltip>
-				<Tooltip
-					content={$i18n.t(
-						'Understand that updating or changing your embedding model requires reset of the vector database and re-import of all documents. You have been warned!'
-					)}
-				>
 					<button
-						class=" self-center text-xs p-1 px-3 bg-gray-100 dark:bg-gray-800 dark:hover:bg-gray-700 rounded flex flex-row space-x-1 items-center {loading1
+						class="px-2.5 bg-gray-100 hover:bg-gray-200 text-gray-800 dark:bg-gray-850 dark:hover:bg-gray-800 dark:text-gray-100 rounded-lg transition"
-							? ' cursor-not-allowed'
-							: ''}"
 						on:click={() => {
-							embeddingModelUpdateHandler(embeddingModel);
+							embeddingModelUpdateHandler();
 						}}
-						type="button"
+						disabled={updateEmbeddingModelLoading}
-						disabled={loading1}
 					>
-						<div class="self-center font-medium">{$i18n.t('Update')}</div>
+						{#if updateEmbeddingModelLoading}
+							<div class="self-center">
-						{#if loading1}
-							<div class="ml-3 self-center">
 								<svg
-									class=" w-3 h-3"
+									class=" w-4 h-4"
 									viewBox="0 0 24 24"
 									fill="currentColor"
 									xmlns="http://www.w3.org/2000/svg"
@@ -234,13 +215,31 @@
 									/></svg
 								>
 							</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>
-				</Tooltip>
 				</div>
+				<div class="mt-2 mb-1 text-xs text-gray-400 dark:text-gray-500">
+					{$i18n.t(
+						'Warning: If you update or change your embedding model, you will need to re-import all documents.'
+					)}
 				</div>
-		<hr class=" dark:border-gray-700" />
+				<hr class=" dark:border-gray-700 my-3" />
 				<div class=" ">
 					<div class=" text-sm font-medium">{$i18n.t('Chunk Params')}</div>
@@ -262,7 +261,9 @@
 						</div>
 						<div class="flex w-full">
-					<div class=" self-center text-xs font-medium min-w-fit">{$i18n.t('Chunk Overlap')}</div>
+							<div class=" self-center text-xs font-medium min-w-fit">
+								{$i18n.t('Chunk Overlap')}
+							</div>
 							<div class="self-center p-3">
 								<input
@@ -440,4 +441,6 @@
 					{$i18n.t('Save')}
 				</button>
 			</div>
+		</div>
+	</div>
 </form>