refac

Dockerfile

@@ -74,6 +74,10 @@ ENV RAG_EMBEDDING_MODEL="$USE_EMBEDDING_MODEL_DOCKER" \
 
 ## Hugging Face download cache ##
 ENV HF_HOME="/app/backend/data/cache/embedding/models"
+
+## Torch Extensions ##
+ENV TORCH_EXTENSIONS_DIR="/app/backend/data/cache/torch_extensions"
+
 #### Other models ##########################################################
 
 WORKDIR /app/backend

backend/open_webui/apps/rag/main.py

@@ -206,18 +206,6 @@ def update_reranking_model(
                     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),