Update main.py

main.py

@@ -61,7 +61,7 @@ index.train(vectors)
 index.add(vectors)
 
 
-def preprocess(query: str, k: int) -> tuple[str, dict]:
+def preprocess(query: str, k: int) -> tuple[str, str]:
     """
     Searches the dataset for the top k most relevant papers to the query and returns a prompt and references
     Args:
@@ -78,14 +78,14 @@ def preprocess(query: str, k: int) -> tuple[str, dict]:
     prompt = (
         "You are an AI assistant who delights in helping people learn about research from the Design Research Collective, which is a research lab at Carnegie Mellon University led by Professor Chris McComb. "
         "Your main task is to provide an ANSWER to the USER_QUERY based on the RESEARCH_ABSTRACTS. "
-        "The RESEARCH_ABSTRACTS are provided in the `.bibtex` format. Your ANSWER should contain citations to the RESEARCH_ABSTRACTS using the LaTeX `\\cite` command (e.g.,: '\\cite{mccomb2015}'). "
+        "The RESEARCH_ABSTRACTS are provided in the `.bibtex` format. Your ANSWER should contain citations to the RESEARCH_ABSTRACTS using (AUTHOR, YEAR) format. "
         "DO NOT list references at the end of the answer.\n\n"
         "RESEARCH_ABSTRACTS:\n```bibtex\n{{ABSTRACTS_GO_HERE}}\n```\n\n"
         "USER_GUERY:\n{{QUERY_GOES_HERE}}\n\n"
         "ANSWER:\n"
     )
 
-    id_to_url = {}
+    references = []
     research_abstracts = ""
 
     for i in range(k):
@@ -106,28 +106,14 @@ def preprocess(query: str, k: int) -> tuple[str, dict]:
         first_authors_last_name = last_names[0]
 
         research_abstracts += top_five["bibtex"].values[i] + "\n"
-        id_to_url[top_five["bib_dict"].values[i]["ID"]] = f"<a href=\"{url}\">[{first_authors_last_name} {year}]</a>"
+        references.append(f"<a href=\"{url}\">{first_authors_last_name} {year}</a>")
         
     prompt = prompt.replace("{{ABSTRACTS_GO_HERE}}", research_abstracts)
     prompt = prompt.replace("{{QUERY_GOES_HERE}}", query)
 
     print(prompt)
     
-    return prompt, id_to_url
-    
-def postprocess(response: str, bypass_from_preprocessing: dict) -> str:
-    """
-    Applies a postprocessing step to the LLM's response before the user receives it
-    Args:
-        response (str): The LLM's response
-        bypass_from_preprocessing (str): The bypass variable from the preprocessing step
-    Returns:
-        str: The postprocessed response
-    """
-    for key in bypass_from_preprocessing.keys():
-        response = response.replace("\\cite{"+key+"}", bypass_from_preprocessing[key])
-        
-    return response
+    return prompt, "; ".join(references)
 
 
 @spaces.GPU
@@ -169,10 +155,7 @@ def reply(message: str, history: list[str]) -> str:
             time.sleep(0.01)
             yield partial_message
 
-    for char in ["\n\n- "]+list("\n- ".join([key+": "+bypass[key] for key in bypass.keys()])):
-        partial_message +=char
-        time.sleep(0.005)
-        yield partial_message
+    yield partial_message + bypass