Update main.py

main.py

@@ -36,6 +36,18 @@ abstract_is_null = [
 data = data[~pandas.Series(abstract_is_null)]
 data.reset_index(inplace=True)
 
+# Load the model for later use in embeddings
+model = sentence_transformers.SentenceTransformer(EMBEDDING_MODEL_NAME)
+
+# Create an LLM pipeline that we can send queries to
+tokenizer = transformers.AutoTokenizer.from_pretrained(LLM_MODEL_NAME)
+streamer = transformers.TextIteratorStreamer(
+    tokenizer, skip_prompt=True, skip_special_tokens=True
+)
+chatmodel = transformers.AutoModelForCausalLM.from_pretrained(
+    LLM_MODEL_NAME, torch_dtype="auto", device_map="auto"
+)
+
 # Create a FAISS index for fast similarity search
 metric = faiss.METRIC_INNER_PRODUCT
 vectors = numpy.stack(data["embedding"].tolist(), axis=0)
@@ -45,34 +57,35 @@ faiss.normalize_L2(vectors)
 index.train(vectors)
 index.add(vectors)
 
-# Load the model for later use in embeddings
-model = sentence_transformers.SentenceTransformer(EMBEDDING_MODEL_NAME)
 
-
-def search(query: str, k: int) -> tuple[str, str]:
+def preprocess(query: str, k: int) -> tuple[str, str]:
     """
-    Searches the dataset for the top k most relevant papers to the query
+    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 search results and references
+        tuple[str, str]: A tuple containing the prompt and references
     """
-    query = numpy.expand_dims(model.encode(query), axis=0)
-    faiss.normalize_L2(query)
-    D, I = index.search(query, k)
+    encoded_query = numpy.expand_dims(model.encode(query), axis=0)
+    faiss.normalize_L2(encoded_query)
+    D, I = index.search(encoded_query, k)
     top_five = data.loc[I[0]]
 
-    search_results = (
+    prompt = (
         "You are an AI assistant who delights in helping people learn about research from the Design "
-        "Research Collective. Your main task is to provide an ANSWER to the USER_QUERY based on the RESEARCH_ABSTRACTS.\n\n"
-        "RESEARCH_ABSTRACTS: \n"
+        "Research Collective. Your main task is to provide an ANSWER to the USER_QUERY based on the "
+        "RESEARCH_ABSTRACTS.\n\n"
+        "RESEARCH_ABSTRACTS:\n{{ABSTRACTS_GO_HERE}}\n\n"
+        "USER_GUERY:\n{{QUERY_GOES_HERE}}\n\n"
+        "ANSWER:\n"
     )
 
     references = "\n\n## References\n\n"
+    research_abstracts = ""
 
     for i in range(k):
-        search_results += top_five["bib_dict"].values[i]["abstract"] + "\n"
+        research_abstracts += top_five["bib_dict"].values[i]["abstract"] + "\n"
         references += (
             str(i + 1)
             + ". "
@@ -93,36 +106,12 @@ def search(query: str, k: int) -> tuple[str, str]:
             + top_five["author_pub_id"].values[i]
             + ").\n"
         )
+        
+    prompt = prompt.replace("{{ABSTRACTS_GO_HERE}}", research_abstracts)
+    prompt = prompt.replace("{{QUERY_GOES_HERE}}", query)
 
-    search_results += (
-        "\n\nUSER_QUERY: "
-    )
-
-    return search_results, references
-
-
-# Create an LLM pipeline that we can send queries to
-tokenizer = transformers.AutoTokenizer.from_pretrained(LLM_MODEL_NAME)
-streamer = transformers.TextIteratorStreamer(
-    tokenizer, skip_prompt=True, skip_special_tokens=True
-)
-chatmodel = transformers.AutoModelForCausalLM.from_pretrained(
-    LLM_MODEL_NAME, torch_dtype="auto", device_map="auto"
-)
-
-
-def preprocess(message: str) -> tuple[str, str]:
-    """
-    Applies a preprocessing step to the user's message before the LLM receives it
-    Args:
-        message (str): The user's message
-    Returns:
-        tuple[str, str]: A tuple containing the preprocessed message and a bypass variable
-    """
-    block_search_results, formatted_search_results = search(message, 5)
-    return block_search_results + message + "\nANSWER: ", formatted_search_results
-
-
+    return prompt, references
+    
 def postprocess(response: str, bypass_from_preprocessing: str) -> str:
     """
     Applies a postprocessing step to the LLM's response before the user receives it
@@ -147,7 +136,7 @@ def reply(message: str, history: list[str]) -> str:
     """
 
     # Apply preprocessing
-    message, bypass = preprocess(message)
+    message, bypass = preprocess(message, 5)
 
     # This is some handling that is applied to the history variable to put it in a good format
     history_transformer_format = [