Update app.py

app.py

@@ -37,19 +37,20 @@ def load_doc(list_file_path, chunk_size, chunk_overlap):
         pages.extend(loader.load())
     text_splitter = RecursiveCharacterTextSplitter(
         chunk_size=chunk_size, 
-        chunk_overlap=chunk_overlap
-    )
+        chunk_overlap=chunk_overlap)
     doc_splits = text_splitter.split_documents(pages)
     return doc_splits
 
+# Create vector database
 def create_db(splits, collection_name):
-    embedding = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", device="cuda")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    embedding = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", device=device)
     new_client = chromadb.EphemeralClient()
     vectordb = Chroma.from_documents(
         documents=splits,
         embedding=embedding,
         client=new_client,
-        collection_name=collection_name
+        collection_name=collection_name,
     )
     return vectordb