UniqAI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2024-11-16 05:24:02 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy J. Baek 2024-09-19 22:17:32 +02:00
parent 4e43663448
commit 9be73ea94a
2 changed files with 13 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
Dockerfile
View File

@ -76,7 +76,7 @@ ENV RAG_EMBEDDING_MODEL="$USE_EMBEDDING_MODEL_DOCKER" \
ENV HF_HOME="/app/backend/data/cache/embedding/models" ENV HF_HOME="/app/backend/data/cache/embedding/models"
## Torch Extensions ## ## Torch Extensions ##
ENV TORCH_EXTENSIONS_DIR="/.cache/torch_extensions" # ENV TORCH_EXTENSIONS_DIR="/.cache/torch_extensions"
#### Other models ########################################################## #### Other models ##########################################################

12
backend/open_webui/apps/rag/main.py
View File

@ -206,6 +206,18 @@ def update_reranking_model(
print("ColBERT: Loading model", name) print("ColBERT: Loading model", name)
self.device = "cuda" if torch.cuda.is_available() else "cpu" self.device = "cuda" if torch.cuda.is_available() else "cpu"
if DOCKER:
# This is a workaround for the issue with the docker container
# where the torch extension is not loaded properly
# and the following error is thrown:
# /root/.cache/torch_extensions/py311_cpu/segmented_maxsim_cpp/segmented_maxsim_cpp.so: cannot open shared object file: No such file or directory
torch_extensions = "/root/.cache/torch_extensions/py311_cpu"
try:
shutil.rmtree(torch_extensions)
except:
pass
self.ckpt = Checkpoint( self.ckpt = Checkpoint(
name, name,
colbert_config=ColBERTConfig(model_name=name), colbert_config=ColBERTConfig(model_name=name),