Update auditqa/doc_process.py

auditqa/doc_process.py

@@ -2,9 +2,11 @@ import glob
 import os
 from langchain.text_splitter import RecursiveCharacterTextSplitter, SentenceTransformersTokenTextSplitter
 from transformers import AutoTokenizer
+from torch import cuda
 from langchain_community.document_loaders import PyMuPDFLoader
 from langchain_community.embeddings import HuggingFaceEmbeddings, HuggingFaceInferenceAPIEmbeddings
 from langchain_community.vectorstores import Qdrant
+device = 'cuda' if cuda.is_available() else 'cpu'
 #from dotenv import load_dotenv
 #load_dotenv()
 
@@ -55,7 +57,7 @@ def process_pdf():
 
 
     embeddings = HuggingFaceEmbeddings(
-        model_kwargs = {'device': 'cpu'},
+        model_kwargs = {'device': device},
         encode_kwargs = {'normalize_embeddings': True},
         model_name="BAAI/bge-small-en-v1.5"
     )