refac

Dockerfile

@@ -157,5 +157,6 @@ USER $UID:$GID
 
 ARG BUILD_HASH
 ENV WEBUI_BUILD_VERSION=${BUILD_HASH}
+ENV DOCKER true
 
 CMD [ "bash", "start.sh"]

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),

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")