Update app.py

app.py

@@ -1,5 +1,4 @@
 import threading  # to allow streaming response
-import time  # to pave the deliver of the message
 
 import gradio  # for the interface
 import spaces  # for GPU
@@ -76,18 +75,16 @@ chatmodel = transformers.AutoModelForCausalLM.from_pretrained(
 )
 
 
-def preprocess(query: str, k: int) -> tuple[str, str]:
+def preprocess(query: str, k: int) -> str:
     """
     Searches the dataset for the top k most relevant papers to the query and returns a prompt and references
     Args:
         query (str): The user's query
         k (int): The number of results to return
     Returns:
-        tuple[str, str]: A tuple containing the prompt and references
+        str: The prompt to be used for the AI
     """
-    documents = publication_vectorstore.search(
-        query, k=PUBLICATIONS_TO_RETRIEVE, search_type="similarity"
-    )
+    documents = publication_vectorstore.search(query, k=k, search_type="similarity")
 
     prompt = (
         "You are an AI assistant who delights in helping people learn about research. "
@@ -109,7 +106,7 @@ def preprocess(query: str, k: int) -> tuple[str, str]:
 
     print(prompt)
 
-    return prompt, ""
+    return prompt
 
 
 @spaces.GPU
@@ -124,7 +121,7 @@ def reply(message: str, history: list[str]) -> str:
     """
 
     # Apply preprocessing
-    message, bypass = preprocess(message, PUBLICATIONS_TO_RETRIEVE)
+    message = preprocess(message, PUBLICATIONS_TO_RETRIEVE)
 
     # This is some handling that is applied to the history variable to put it in a good format
     history_transformer_format = [
@@ -150,6 +147,7 @@ def reply(message: str, history: list[str]) -> str:
             partial_message += new_token
             yield partial_message
 
+
 # Example queries
 EXAMPLE_QUERIES = [
     "What is multi-material 3D printing?",
@@ -162,7 +160,6 @@ EXAMPLE_QUERIES = [
     "What are the benefits and limitations of using polymers in 3D printing?",
     "Tell me about the environmental impacts of additive manufacturing.",
     "What are the primary limitations of current 3D printing technologies?",
-    "What future trends are expected in the field of additive manufacturing?",
     "How are researchers improving the speed of 3D printing processes?",
     "What are the best practices for managing post-processing in additive manufacturing?",
 ]