Fix Index root is not used correctly in server.py

stanford-futuredata · Oct 26, 2023 · 7c7113b · 7c7113b
1 parent cd8582f
commit 7c7113b
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/colbert/searcher.py b/colbert/searcher.py
@@ -20,13 +20,14 @@
 
 
 class Searcher:
-    def __init__(self, index, checkpoint=None, collection=None, config=None):
+    def __init__(self, index, checkpoint=None, collection=None, config=None, index_root=None):
         print_memory_stats()
 
         initial_config = ColBERTConfig.from_existing(config, Run().config)
 
         default_index_root = initial_config.index_root_
-        self.index = os.path.join(default_index_root, index)
+        index_root = index_root if index_root else default_index_root
+        self.index = os.path.join(index_root, index)
         self.index_config = ColBERTConfig.load_from_index(self.index)
 
         self.checkpoint = checkpoint or self.index_config.checkpoint

diff --git a/server.py b/server.py
@@ -13,7 +13,7 @@
 INDEX_ROOT = os.getenv("INDEX_ROOT")
 app = Flask(__name__)
 
-searcher = Searcher(index=f"{INDEX_ROOT}/{INDEX_NAME}")
+searcher = Searcher(index=INDEX_NAME, index_root=INDEX_ROOT)
 counter = {"api" : 0}
 
 @lru_cache(maxsize=1000000)