Update app.py

app.py

@@ -51,20 +51,32 @@ except PermissionError as e:
 document_store = {}
 
 # Check if FAISS index file exists, otherwise initialize it
-if os.path.exists(faiss_index_file):
-    try:
-        index = faiss.read_index(faiss_index_file)
-        if index.ntotal > 0:
-            print(f"FAISS index loaded with {index.ntotal} vectors.")
-            index.reset()  # Resetting the index if non-zero entries
-            index = faiss.IndexIDMap(faiss.IndexFlatL2(384))  # Reinitialize the index
-        else:
-            index = faiss.IndexIDMap(faiss.IndexFlatL2(384))  # Initialize with flat L2 distance
-    except Exception as e:
-        print(f"Error loading FAISS index: {e}, reinitializing.")
-        index = faiss.IndexIDMap(faiss.IndexFlatL2(384))  # Initialize if reading fails
-else:
-    index = faiss.IndexIDMap(faiss.IndexFlatL2(384))  # Initialize if file doesn't exist
+def initialize_faiss():
+    if os.path.exists(faiss_index_file):
+        try:
+            index = faiss.read_index(faiss_index_file)
+            if index.ntotal > 0:
+                print(f"FAISS index loaded with {index.ntotal} vectors.")
+                index.reset()  # Resetting the index if non-zero entries
+                index = faiss.IndexIDMap(faiss.IndexFlatL2(384))  # Reinitialize the index
+            else:
+                index = faiss.IndexIDMap(faiss.IndexFlatL2(384))  # Initialize with flat L2 distance
+        except Exception as e:
+            print(f"Error loading FAISS index: {e}, reinitializing.")
+            index = faiss.IndexIDMap(faiss.IndexFlatL2(384))  # Initialize if reading fails
+    else:
+        index = faiss.IndexIDMap(faiss.IndexFlatL2(384))  # Initialize if file doesn't exist
+
+    # Check if GPU is available and move index to GPU if possible
+    if torch.cuda.is_available():
+        print("CUDA is available, moving FAISS index to GPU.")
+        index = faiss.index_cpu_to_all_gpus(index)
+    else:
+        print("CUDA not available, using CPU.")
+    
+    return index
+
+index = initialize_faiss()
 
 # Function to upload document
 def upload_document(file_path, embed_model):