Update app.py

app.py

@@ -5,8 +5,8 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import gradio as gr
 import json
 import faiss
-import spaces
 import numpy as np
+import spaces
 # Ensure you have GPU support
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
@@ -32,7 +32,6 @@ 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)
@@ -40,14 +39,23 @@ 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
         relevant_doc = summarizer(relevant_doc, max_length=150, min_length=50, do_sample=False)[0]['summary_text']
+    
     input_text = f"Document: {relevant_doc}\n\nQuestion: {query}\n\nAnswer:"
     inputs = llama_tokenizer(input_text, return_tensors="pt").to(device)
-    outputs = llama_model.generate(inputs["input_ids"], max_length=150)
+    
+    # Set pad_token_id to eos_token_id to avoid the warning
+    pad_token_id = llama_tokenizer.eos_token_id
+    outputs = llama_model.generate(
+        inputs["input_ids"], 
+        attention_mask=inputs["attention_mask"], 
+        max_length=150, 
+        pad_token_id=pad_token_id
+    )
+    
     response = llama_tokenizer.decode(outputs[0], skip_special_tokens=True)
     return response