Update app.py (#3)

- Update app.py (3b889f09a94995033c8a5d8899a6b88e7f3b0347)


Co-authored-by: Aryan Wadhawan <aryanxxvii@users.noreply.huggingface.co>

app.py

@@ -38,56 +38,57 @@ embeddings = SentenceTransformerEmbeddings(
     device=device
 )
 
-# Define vectorstore paths consistently
-VECTORSTORE_DIR = "vectorstore/db_faiss"
-vectorstore_path = pathlib.Path(VECTORSTORE_DIR)
-
-# Create vectorstore directory if it doesn't exist
-vectorstore_path.mkdir(parents=True, exist_ok=True)
-
-if not (vectorstore_path / "index.faiss").exists():
-    print("Creating new vectorstore...")
-    # Load and split the PDF
-    loader = PyPDFLoader("medical.pdf")
-    documents = loader.load()
+def create_vectorstore():
+    # Define vectorstore paths consistently
+    VECTORSTORE_DIR = "vectorstore/db_faiss"
+    vectorstore_path = pathlib.Path(VECTORSTORE_DIR)
     
-    # Split documents into chunks
-    text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=2000,
-        chunk_overlap=100,
-        length_function=len,
-    )
-    texts = text_splitter.split_documents(documents)
+    # Create vectorstore directory if it doesn't exist
+    vectorstore_path.mkdir(parents=True, exist_ok=True)
     
-    # Create and save the vectorstore
-    vectorstore = FAISS.from_documents(texts, embeddings)
-    
-    # If CUDA is available, convert index to GPU
-    if device == "cuda":
-        res = faiss.StandardGpuResources()  # Initialize GPU resources
-        index = vectorstore.index
-        gpu_index = faiss.index_cpu_to_gpu(res, 0, index)  # Move to GPU
-        vectorstore.index = gpu_index
-    
-    # Save the vectorstore
-    vectorstore.save_local(VECTORSTORE_DIR)
-    print("Vectorstore created and saved successfully.")
-else:
-    print("Loading existing vectorstore...")
-    # Load existing vectorstore
-    vectorstore = FAISS.load_local(
-        folder_path=VECTORSTORE_DIR,
-        embeddings=embeddings,
-        allow_dangerous_deserialization=True
-    )
-    
-    # If CUDA is available, convert loaded index to GPU
-    if device == "cuda":
-        res = faiss.StandardGpuResources()  # Initialize GPU resources
-        index = vectorstore.index
-        gpu_index = faiss.index_cpu_to_gpu(res, 0, index)  # Move to GPU
-        vectorstore.index = gpu_index
-    print("Vectorstore loaded successfully.")
+    if not (vectorstore_path / "index.faiss").exists():
+        print("Creating new vectorstore...")
+        # Load and split the PDF
+        loader = PyPDFLoader("medical.pdf")
+        documents = loader.load()
+        
+        # Split documents into chunks
+        text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=2000,
+            chunk_overlap=100,
+            length_function=len,
+        )
+        texts = text_splitter.split_documents(documents)
+        
+        # Create and save the vectorstore
+        vectorstore = FAISS.from_documents(texts, embeddings)
+        
+        # If CUDA is available, convert index to GPU
+        if device == "cuda":
+            res = faiss.StandardGpuResources()  # Initialize GPU resources
+            index = vectorstore.index
+            gpu_index = faiss.index_cpu_to_gpu(res, 0, index)  # Move to GPU
+            vectorstore.index = gpu_index
+        
+        # Save the vectorstore
+        vectorstore.save_local(VECTORSTORE_DIR)
+        print("Vectorstore created and saved successfully.")
+    else:
+        print("Loading existing vectorstore...")
+        # Load existing vectorstore
+        vectorstore = FAISS.load_local(
+            folder_path=VECTORSTORE_DIR,
+            embeddings=embeddings,
+            allow_dangerous_deserialization=True
+        )
+        
+        # If CUDA is available, convert loaded index to GPU
+        if device == "cuda":
+            res = faiss.StandardGpuResources()  # Initialize GPU resources
+            index = vectorstore.index
+            gpu_index = faiss.index_cpu_to_gpu(res, 0, index)  # Move to GPU
+            vectorstore.index = gpu_index
+        print("Vectorstore loaded successfully.")
 
 def get_relevant_context(query):
     try: