Update app.py

app.py

@@ -32,6 +32,7 @@ llama_model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2").to(d
 summarizer = pipeline("summarization", model="facebook/bart-large-cnn", device=0 if device == 'cuda' else -1)
 
 # Define the function to find the most relevant document using FAISS
+@spaces.GPU(duration=120)
 def retrieve_relevant_doc(query):
     query_embedding = sentence_model.encode(query, convert_to_tensor=False)
     _, indices = index.search(np.array([query_embedding]), k=1)
@@ -39,6 +40,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)
     if len(relevant_doc) > 512:  # Truncate long documents