commit to load hf vectorstore index

code/modules/config/config.yml

@@ -3,11 +3,13 @@ log_chunk_dir: '../storage/logs/chunks' # str
 device: 'cpu' # str [cuda, cpu]
 
 vectorstore:
+  load_from_HF: True # bool
+  HF_path: "XThomasBU/Colbert_Index" # str
   embedd_files: False # bool
   data_path: '../storage/data' # str
   url_file_path: '../storage/data/urls.txt' # str
   expand_urls: True # bool
-  db_option : 'FAISS' # str [FAISS, Chroma, RAGatouille, RAPTOR]
+  db_option : 'RAGatouille' # str [FAISS, Chroma, RAGatouille, RAPTOR]
   db_path : '../vectorstores' # str
   model : 'sentence-transformers/all-MiniLM-L6-v2' # str [sentence-transformers/all-MiniLM-L6-v2, text-embedding-ada-002']
   search_top_k : 3 # int

code/modules/vectorstore/store_manager.py

@@ -143,6 +143,14 @@ class VectorStoreManager:
         self.logger.info("Loaded database")
         return self.loaded_vector_db
 
+    def load_from_HF(self):
+        start_time = time.time()  # Start time for loading database
+        self.vector_db._load_from_HF()
+        end_time = time.time()
+        self.logger.info(
+            f"Time taken to load database from Hugging Face: {end_time - start_time} seconds"
+        )
+
 
 if __name__ == "__main__":
     import yaml
@@ -152,7 +160,10 @@ if __name__ == "__main__":
     print(config)
     print(f"Trying to create database with config: {config}")
     vector_db = VectorStoreManager(config)
-    vector_db.create_database()
+    if config["vectorstore"]["load_from_HF"] and "HF_path" in config["vectorstore"]:
+        vector_db.load_from_HF()
+    else:
+        vector_db.create_database()
     print("Created database")
 
     print(f"Trying to load the database")

code/modules/vectorstore/vectorstore.py

@@ -2,6 +2,9 @@ from modules.vectorstore.faiss import FaissVectorStore
 from modules.vectorstore.chroma import ChromaVectorStore
 from modules.vectorstore.colbert import ColbertVectorStore
 from modules.vectorstore.raptor import RAPTORVectoreStore
+from huggingface_hub import snapshot_download
+import os
+import shutil
 
 
 class VectorStore:
@@ -50,6 +53,32 @@ class VectorStore:
         else:
             return self.vectorstore.load_database(embedding_model)
 
+    def _load_from_HF(self):
+        # Download the snapshot from Hugging Face Hub
+        # Note: Download goes to the cache directory
+        snapshot_path = snapshot_download(
+            repo_id=self.config["vectorstore"]["HF_path"], repo_type="dataset"
+        )
+
+        # Move the downloaded files to the desired directory
+        target_path = os.path.join(
+            self.config["vectorstore"]["db_path"],
+            "db_" + self.config["vectorstore"]["db_option"],
+        )
+
+        # Create target path if it doesn't exist
+        os.makedirs(target_path, exist_ok=True)
+
+        # move all files and directories from snapshot_path to target_path
+        # target path is used while loading the database
+        for item in os.listdir(snapshot_path):
+            s = os.path.join(snapshot_path, item)
+            d = os.path.join(target_path, item)
+            if os.path.isdir(s):
+                shutil.copytree(s, d, dirs_exist_ok=True)
+            else:
+                shutil.copy2(s, d)
+
     def _as_retriever(self):
         return self.vectorstore.as_retriever()