Update app.py

app.py

@@ -9,8 +9,11 @@ def setup():
     from stream_unzip import stream_unzip
     
     DATASET_URL = os.environ.get("DATASET_URL")
+    INDEX_TYPE  = os.environ.get("INDEX_TYPE")
+    
     if not DATASET_URL:
         raise ValueError("DATASET_URL must be set in the environment")
+        
     extract_dir = "/dev/shm/rpc-vecdb"
     os.makedirs(extract_dir, exist_ok=True)
     response = requests.get(DATASET_URL, stream=True)
@@ -32,12 +35,17 @@ def setup():
     for f in files: print(f)
     print("Index extracted")
 
-    print("Installing NGT...")
-    subprocess.check_call(["bash", "install_ngt.sh"])
-    print("NGT installed")
+    if INDEX_TYPE == "ngt":
+        print("Installing NGT...")
+        subprocess.check_call(["bash", "install_ngt.sh"])
+        print("NGT installed")
 
     print("Loading index...")
-    load_index(extract_dir + "/index")
+    if INDEX_TYPE == "ngt":
+        index_dir = extract_dir + "/index"
+    else:
+        index_dir = extract_dir
+    load_index(index_path=index_dir, idx_type=INDEX_TYPE)
     print("Index loaded")