Update app.py

app.py

@@ -29,6 +29,7 @@ list_llm = [
 ]
 list_llm_simple = [os.path.basename(llm) for llm in list_llm]
 
+@spaces.GPU
 @spaces.GPU
 def load_doc(list_file_path, chunk_size, chunk_overlap):
     loaders = [PyPDFLoader(x) for x in list_file_path]
@@ -36,15 +37,18 @@ def load_doc(list_file_path, chunk_size, chunk_overlap):
     for loader in loaders:
         pages.extend(loader.load())
     text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=chunk_size, 
-        chunk_overlap=chunk_overlap)
+        chunk_size = chunk_size, 
+        chunk_overlap = chunk_overlap)
     doc_splits = text_splitter.split_documents(pages)
     return doc_splits
 
 # Create vector database
 def create_db(splits, collection_name):
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    embedding = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", device=device)
+    # Set CUDA_VISIBLE_DEVICES if GPU is available
+    if torch.cuda.is_available():
+        os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+    
+    embedding = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
     new_client = chromadb.EphemeralClient()
     vectordb = Chroma.from_documents(
         documents=splits,