Update app.py

app.py

@@ -2,12 +2,7 @@ import os
 import tempfile
 import torch
 import gradio as gr
-
-# Patch torch if needed
-if not hasattr(torch, "get_default_device"):
-    def get_default_device():
-        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    torch.get_default_device = get_default_device
+import spaces  # Required for GPU-enabled Spaces
 
 from pinecone import Pinecone
 from langchain_pinecone import PineconeVectorStore
@@ -17,21 +12,25 @@ from langchain.chains import RetrievalQAWithSourcesChain
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain.embeddings import HuggingFaceEmbeddings
 
-# Set device
+# Set device for embeddings
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Initialize Pinecone
 pc = Pinecone(api_key=os.getenv("PINECONE_API_KEY"))
 INDEX_NAME = "ragreader"
 
-# Embedding initialization
-embeddings = HuggingFaceEmbeddings(
-    model_name="BAAI/bge-large-en-v1.5",
-    model_kwargs={'device': device}
-)
+# GPU-decorated function to load HuggingFace embeddings on GPU
+@spaces.GPU
+def init_embeddings():
+    return HuggingFaceEmbeddings(
+        model_name="BAAI/bge-large-en-v1.5",
+        model_kwargs={"device": device}
+    )
 
+embeddings = init_embeddings()
 
-# Document processing pipeline
+# GPU-decorated document processing function
+@spaces.GPU
 def process_documents(files):
     docs = []
     for file in files:
@@ -54,7 +53,7 @@ def process_documents(files):
     )
     return "Documents processed and stored."
 
-# Initialize QA chain
+# Initialize the RetrievalQA chain (no GPU decoration needed here)
 def init_qa_chain():
     llm = ChatMistralAI(
         model="mistral-tiny",
@@ -76,7 +75,7 @@ def init_qa_chain():
 
 qa_chain = None
 
-# Gradio interface
+# Build the Gradio UI
 with gr.Blocks() as demo:
     gr.Markdown("## RAG Chatbot - PDF Reader")