Update app.py

app.py

@@ -4,6 +4,7 @@ from sentence_transformers import SentenceTransformer, util
 import gradio as gr
 import json
 from transformers import AutoTokenizer, AutoModelForCausalLM
+import spaces
 
 # Ensure you have GPU support
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -23,6 +24,7 @@ llama_tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
 llama_model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2").to(device)
 
 # Define the function to find the most relevant document
+@spaces.GPU(duration=120)
 def retrieve_relevant_doc(query):
     query_embedding = model.encode(query, convert_to_tensor=True, device=device)
     similarities = util.pytorch_cos_sim(query_embedding, embeddings)[0]
@@ -30,6 +32,7 @@ def retrieve_relevant_doc(query):
     return df.iloc[best_match_idx]['Abstract']
 
 # Define the function to generate a response
+@spaces.GPU(duration=120)
 def generate_response(query):
     relevant_doc = retrieve_relevant_doc(query)
     input_text = f"Document: {relevant_doc}\n\nQuestion: {query}\n\nAnswer:"