diff --git a/Dockerfile b/Dockerfile index 550708ba8..795a1141a 100644 --- a/Dockerfile +++ b/Dockerfile @@ -157,5 +157,6 @@ USER $UID:$GID ARG BUILD_HASH ENV WEBUI_BUILD_VERSION=${BUILD_HASH} +ENV DOCKER true CMD [ "bash", "start.sh"] diff --git a/backend/open_webui/apps/rag/main.py b/backend/open_webui/apps/rag/main.py index 5b7452af5..3ba52ae7f 100644 --- a/backend/open_webui/apps/rag/main.py +++ b/backend/open_webui/apps/rag/main.py @@ -90,7 +90,7 @@ from open_webui.config import ( AppConfig, ) from open_webui.constants import ERROR_MESSAGES -from open_webui.env import SRC_LOG_LEVELS, DEVICE_TYPE +from open_webui.env import SRC_LOG_LEVELS, DEVICE_TYPE, DOCKER from open_webui.utils.misc import ( calculate_sha256, calculate_sha256_string, @@ -205,6 +205,19 @@ def update_reranking_model( def __init__(self, name) -> None: print("ColBERT: Loading model", name) 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 + + lock_file = "/root/.cache/torch_extensions/py311_cpu/segmented_maxsim_cpp/lock" + if os.path.exists(lock_file): + os.remove(lock_file) + print("ColBERT: Removed lock file") + + self.ckpt = Checkpoint( name, colbert_config=ColBERTConfig(model_name=name), diff --git a/backend/open_webui/env.py b/backend/open_webui/env.py index cf45e725c..89422e57b 100644 --- a/backend/open_webui/env.py +++ b/backend/open_webui/env.py @@ -31,6 +31,7 @@ try: except ImportError: print("dotenv not installed, skipping...") +DOCKER = os.environ.get("DOCKER", "False").lower() == "true" # device type embedding models - "cpu" (default), "cuda" (nvidia gpu required) or "mps" (apple silicon) - choosing this right can lead to better performance USE_CUDA = os.environ.get("USE_CUDA_DOCKER", "false")