Update app.py

app.py

@@ -32,6 +32,7 @@ llama_model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-h
 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
@@ -52,7 +54,7 @@ def generate_response(query):
 # Create a Gradio interface
 iface = gr.Interface(
     fn=generate_response,
-    inputs=gr.inputs.Textbox(lines=2, placeholder="Enter your query here..."),
+    inputs=gr.Textbox(lines=2, placeholder="Enter your query here..."),
     outputs="text",
     title="RAG Chatbot",
     description="This chatbot retrieves relevant documents based on your query and generates responses using LLaMA."