Update app.py

app.py

@@ -83,9 +83,9 @@ def initialize_faiss():
 def save_faiss_index(index):
     try:
         if torch.cuda.is_available():
-            print("Moving FAISS index back to CPU before saving.")
-            index = faiss.index_gpu_to_cpu(index)  # Move back to CPU before saving
-
+            print("Moving FAISS index back to CPU before saving.") 
+            res = faiss.StandardGpuResources()  # Allocate GPU resources
+            index = faiss.index_cpu_to_gpu(res, 0, index)  # Move to GPU 0
         print(f"Saving FAISS index to {faiss_index_file}.")
         faiss.write_index(index, faiss_index_file)
         print(f"FAISS index successfully saved to {faiss_index_file}.")
@@ -101,7 +101,7 @@ save_faiss_index(index)
 # Load document store and populate FAISS index
 knowledgebase_file = os.path.join(UPLOAD_DIR, "knowledge_text.txt")  # Ensure this path is correct
 
-def load_document_store():
+def add_document_store_to_index():
     """Loads knowledgebase.txt into a dictionary where FAISS IDs map to text and embeddings"""
     global document_store
     document_store = {}  # Reset document store
@@ -121,15 +121,36 @@ def load_document_store():
     else:
         print("Error: knowledgebase.txt not found!")
     
-    # Generate embeddings for all documents
-    embeddings = bertmodel.encode(all_texts)
-    embeddings = embeddings.astype("float32")
-    
+    # Generate embeddings for all documents 
+    embeddings = bertmodel.encode(all_texts, batch_size=32, convert_to_numpy=True).astype("float32")
+
     # Add embeddings to FAISS index
     index.add_with_ids(embeddings, np.array(list(document_store.keys()), dtype=np.int64))
     print(f"Added {len(all_texts)} document embeddings to FAISS index.")
  
+add_document_store_to_index()
 
+def load_document_store():
+    """Loads knowledgebase.txt into a dictionary where FAISS IDs map to text and embeddings"""
+    global document_store
+    document_store = {}  # Reset document store
+    all_texts = []
+    
+    if os.path.exists(knowledgebase_file):
+        with open(knowledgebase_file, "r", encoding="utf-8") as f:
+            lines = f.readlines()
+        
+        for i, line in enumerate(lines):
+            text = line.strip()
+            if text:
+                document_store[i] = {"text": text}  # Store text mapped to FAISS ID
+                all_texts.append(text)  # Collect all texts for embedding
+
+        print(f"Loaded {len(document_store)} documents into document_store.")
+    else:
+        print("Error: knowledgebase.txt not found!")
+     
+ 
 # Function to upload document
 def upload_document(file_path, embed_model):
     try:
@@ -315,10 +336,12 @@ def generate_response(user_input, model_id):
        
         # Append chat history
         func_caller = []
-         
+          
         query_vector = bertmodel.encode(user_input).reshape(1, -1).astype("float32")
         D, I = index.search(query_vector, 1)
 
+        load_document_store()
+
         # Retrieve document
         retrieved_id = I[0][0]
         retrieved_knowledge = (