Update app.py

app.py

@@ -64,6 +64,7 @@
 
 
 
+
 import gradio as gr
 import torch
 from transformers import RagRetriever, RagSequenceForGeneration, AutoTokenizer
@@ -71,7 +72,7 @@ from transformers import RagRetriever, RagSequenceForGeneration, AutoTokenizer
 """
 For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
 """
-device = torch.device("cuda" if torch.cuda.is_available() else "CPU")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 def strip_title(title):
     if title.startswith('"'):
@@ -113,6 +114,7 @@ def retrieved_info(rag_model, query):
             retrieved_context.append(f"{title}: {text}")
 
     answer = retrieved_context
+    return answer
 
 
 
@@ -121,33 +123,33 @@ def respond(
     message,
     history: list[tuple[str, str]],
     system_message,
-    max_tokens = None,
-    temperature = None,
-    top_p = None,
+    max_tokens ,
+    temperature,
+    top_p,
 ):
     # Load model
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    
+
     dataset_path = "./sample/my_knowledge_dataset"
     index_path = "./sample/my_knowledge_dataset_hnsw_index.faiss"
-    
+
     tokenizer = AutoTokenizer.from_pretrained("facebook/rag-sequence-nq")
     retriever = RagRetriever.from_pretrained("facebook/rag-sequence-nq", index_name="custom",
                                              passages_path = dataset_path,
                                              index_path = index_path,
-                                             n_docs = 1)
+                                             n_docs = 5)
     rag_model = RagSequenceForGeneration.from_pretrained('facebook/rag-sequence-nq', retriever=retriever)
     rag_model.retriever.init_retrieval()
     rag_model.to(device)
 
     if message:  # If there's a user query
         response = retrieved_info(rag_model, message)  # Get the answer from your local FAISS and Q&A model
-        return response
+        return response[0]
 
     # In case no message, return an empty string
     return ""
-    
-    
+
+
 
 """
 For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
@@ -170,4 +172,4 @@ demo = gr.ChatInterface(
 
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch( )