tdpeuter/open-webui
tdpeuter/open-webui
1
Fork 0
forked from open-webui/open-webui
Code Issues Pull requests 1 Projects 1 Releases 2 Packages 1 Activity Actions

refac

Browse source
This commit is contained in:
Timothy J. Baek 2024-04-10 00:33:45 -07:00
parent 48aad65514
commit f4b87ecb23
3 changed files with 243 additions and 236 deletions
Download patch file Download diff file Expand all files Collapse all files

10
backend/apps/rag/main.py
View file

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

8
backend/apps/rag/utils.py
View file

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