import spaces

app.py

@@ -10,6 +10,7 @@ from typing import List
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import ChatPromptTemplate
 import gradio as gr
+import spaces
 from huggingface_hub import InferenceClient
 
 # Configuration
@@ -48,8 +49,16 @@ else:
 dataset = load_dataset('tosin2013/autogen', streaming=True)
 dataset = Dataset.from_list(list(dataset['train']))
 
-# Initialize embeddings with GPU support if available
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# Check GPU availability and initialize embeddings
+if torch.cuda.is_available():
+    print("[INFO] GPU is available, using CUDA")
+    device = "cuda"
+    # Initialize CUDA device
+    torch.cuda.init()
+else:
+    print("[INFO] GPU not available, using CPU")
+    device = "cpu"
+
 embeddings = HuggingFaceEmbeddings(
     model_name="sentence-transformers/all-MiniLM-L6-v2",
     model_kwargs={"device": device}
@@ -65,6 +74,7 @@ text_embeddings = embeddings.embed_documents(texts)
 nn = NearestNeighbors(n_neighbors=5, metric='cosine')
 nn.fit(np.array(text_embeddings))
 
+@spaces.GPU(duration=30)
 def get_relevant_documents(query, k=5):
     """
     Retrieves the k most relevant documents to the query.
@@ -74,6 +84,7 @@ def get_relevant_documents(query, k=5):
     relevant_docs = [texts[i] for i in indices[0]]
     return relevant_docs
 
+@spaces.GPU(duration=120)
 def generate_response(question, history):
     try:
         print(f"\n[LOG] Received question: {question}")
@@ -199,9 +210,10 @@ with gr.Blocks() as demo:
     
     # Event handlers
     submit_btn.click(
-        generate_response,
+        fn=generate_response,
         inputs=[question, chatbot],
-        outputs=[chatbot]
+        outputs=[chatbot],
+        queue=True
     )
     
     clear_btn.click(