Update app.py

app.py

@@ -1,5 +1,6 @@
 import uuid
 import chromadb
+import torch
 from langchain.vectorstores import Chroma
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.retrievers import ContextualCompressionRetriever
@@ -7,8 +8,15 @@ from langchain.retrievers.document_compressors import CrossEncoderReranker
 from langchain_community.cross_encoders import HuggingFaceCrossEncoder
 import gradio as gr
 
+# Set device to GPU if available, else CPU
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
+
 # Initialize embedding model
-embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+embedding_model = HuggingFaceEmbeddings(
+    model_name="sentence-transformers/all-MiniLM-L6-v2",
+    model_kwargs={"device": device}
+)
 
 # Initialize ChromaDB client and collection
 chroma_client = chromadb.PersistentClient(path="./chroma_db")
@@ -19,7 +27,10 @@ vectorstore = Chroma(
 )
 
 # Initialize reranker
-reranker = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-base")
+reranker = HuggingFaceCrossEncoder(
+    model_name="BAAI/bge-reranker-base",
+    model_kwargs={"device": device}
+)
 compressor = CrossEncoderReranker(model=reranker, top_n=5)
 retriever = vectorstore.as_retriever(search_kwargs={"k": 10})  # Retrieve 2k initially
 compression_retriever = ContextualCompressionRetriever(