From 9be73ea94a16afa7e3ff61bd83389283838b8e83 Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Thu, 19 Sep 2024 22:17:32 +0200
Subject: [PATCH] refac

---
 Dockerfile                          |  2 +-
 backend/open_webui/apps/rag/main.py | 12 ++++++++++++
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/Dockerfile b/Dockerfile
index 7389c1ee3..c86f96bf2 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -76,7 +76,7 @@ ENV RAG_EMBEDDING_MODEL="$USE_EMBEDDING_MODEL_DOCKER" \
 ENV HF_HOME="/app/backend/data/cache/embedding/models"
 
 ## Torch Extensions ##
-ENV TORCH_EXTENSIONS_DIR="/.cache/torch_extensions"
+# ENV TORCH_EXTENSIONS_DIR="/.cache/torch_extensions"
 
 #### Other models ##########################################################
 
diff --git a/backend/open_webui/apps/rag/main.py b/backend/open_webui/apps/rag/main.py
index fb3b2fb43..5a228fcea 100644
--- a/backend/open_webui/apps/rag/main.py
+++ b/backend/open_webui/apps/rag/main.py
@@ -206,6 +206,18 @@ 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
+
+                        torch_extensions = "/root/.cache/torch_extensions/py311_cpu"
+                        try:
+                            shutil.rmtree(torch_extensions)
+                        except:
+                            pass
+
                     self.ckpt = Checkpoint(
                         name,
                         colbert_config=ColBERTConfig(model_name=name),