knn_folder

hydra_node/config.py

@@ -146,6 +146,7 @@ def init_config_model():
     config.penultimate = os.getenv('PENULTIMATE', "0")
     config.vae_path = os.getenv('VAE_PATH', None)
     config.module_path = os.getenv('MODULE_PATH', None)
+    config.knn_folder = os.getenv('KNN_FOLDER', None)
     config.prior_path = os.getenv('PRIOR_PATH', None)
     config.default_config = os.getenv('DEFAULT_CONFIG', None)
     config.quality_hack = os.getenv('QUALITY_HACK', "0")

hydra_node/models.py

@@ -707,10 +707,11 @@ class EmbedderModel(nn.Module):
         import faiss
         import pickle
         import requests
+        knn_folder = config.knn_folder 
         self.model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2').cuda()
-        self.index = pickle.load(requests.get("https://f004.backblazeb2.com/file/naipublicbucketxyz/safe/index.pkl", stream='True').raw)
-        self.tag_count = pickle.load(requests.get("https://f004.backblazeb2.com/file/naipublicbucketxyz/safe/all_tags.pkl", stream='True').raw)
-        r = requests.get("https://f004.backblazeb2.com/file/naipublicbucketxyz/safe/knn.index", stream='True')
+        self.index = pickle.load(requests.get(f"{knn_folder}/index.pkl", stream='True').raw)
+        self.tag_count = pickle.load(requests.get(f"{knn_folder}/all_tags.pkl", stream='True').raw)
+        r = requests.get(f"{knn_folder}/knn.index", stream='True')
         with open("knn.index", "wb") as f:
             f.write(r.content)