Update app.py

app.py

@@ -24,28 +24,25 @@ data['embedding'] = data['embedding'].apply(safe_json_loads)
 # Filter out any rows with empty embeddings
 data = data[data['embedding'].apply(lambda x: x.size > 0)]
 
-# Check if the DataFrame is empty after filtering
-if data.empty:
-    raise RuntimeError("No valid embeddings found in the data.")
-
 # Initialize FAISS index
-dimension = len(data['embedding'].iloc[0])
-gpu_available = torch.cuda.is_available()
-
-# Initialize FAISS resources and index
-res = faiss.StandardGpuResources() if gpu_available else None
-index = faiss.IndexFlatL2(dimension)
+dimension = len(data['embedding'][0])
+res = faiss.StandardGpuResources()  # use a single GPU
 
-if gpu_available:
-    index = faiss.index_cpu_to_gpu(res, 0, index)  # move to GPU
+# Check available GPU devices
+num_gpus = faiss.get_num_gpus()
+if num_gpus > 0:
+    gpu_index = faiss.IndexFlatL2(dimension)
+    gpu_index = faiss.index_cpu_to_gpu(res, 0, gpu_index)  # move to GPU
+else:
+    raise RuntimeError("No GPU devices available.")
 
-index.add(np.stack(data['embedding'].values))
+gpu_index.add(np.stack(data['embedding'].values))
 
-# Set the device
-device = torch.device('cuda' if gpu_available else 'cpu')
+# Check if GPU is available
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 # Load QA model
-qa_model = pipeline("question-answering", model="distilbert-base-uncased-distilled-squad", device=0 if gpu_available else -1)
+qa_model = pipeline("question-answering", model="distilbert-base-uncased-distilled-squad", device=0 if torch.cuda.is_available() else -1)
 
 # Load BERT model and tokenizer
 tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
@@ -65,7 +62,7 @@ def retrieve_and_generate(question):
     question_embedding = embed_question(question, model, tokenizer)
     
     # Search in FAISS index
-    _, indices = index.search(question_embedding, k=1)
+    _, indices = gpu_index.search(question_embedding, k=1)
     
     # Retrieve the most relevant document
     relevant_doc = data.iloc[indices[0][0]]