Some updates but mostly formatting

main.py

@@ -1,29 +1,35 @@
-import gradio  # for the interface
-import transformers  # to load an LLM
-import sentence_transformers  # to load an embedding model
+import json  # to work with JSON
+import threading  # for threading
+import time  # for better HCI
+
+import datasets  # to load the dataset
 import faiss  # to create an index
+import gradio  # for the interface
 import numpy  # to work with vectors
 import pandas  # to work with pandas
-import json  # to work with JSON
-import datasets  # to load the dataset
+import sentence_transformers  # to load an embedding model
 import spaces  # for GPU
-import threading  # for threading
-import time   # for better HCI
+import transformers  # to load an LLM
 
 # Constants
-GREETING = "Hi there! I'm an AI agent that uses a [retrieval-augmented generation](https://en.wikipedia.org/wiki/Retrieval-augmented_generation) pipeline to answer questions about research by the Design Research Collective. And the best part is that I always cite my ssources! What can I tell you about today?"
+GREETING = (
+    "Howdy! I'm an AI agent that uses a [retrieval-augmented generation]("
+    "https://en.wikipedia.org/wiki/Retrieval-augmented_generation) pipeline to answer questions about research by the "
+    "Design Research Collective. And the best part is that I always cite my ssources! What can I tell you about today?"
+)
 EMBEDDING_MODEL_NAME = "allenai-specter"
-LLM_MODEL_NAME = "Qwen/Qwen2-7B-Instruct"
+# LLM_MODEL_NAME = "Qwen/Qwen2-7B-Instruct"
+LLM_MODEL_NAME = "Qwen/Qwen2-0.5B-Instruct"
 
 # Load the dataset and convert to pandas
 full_data = datasets.load_dataset("ccm/publications")["train"].to_pandas()
 
 # Filter out any publications without an abstract
-filter = [
+abstract_is_null = [
     '"abstract": null' in json.dumps(bibdict)
     for bibdict in full_data["bib_dict"].values
 ]
-data = full_data[~pandas.Series(filter)]
+data = full_data[~pandas.Series(abstract_is_null)]
 data.reset_index(inplace=True)
 
 # Create a FAISS index for fast similarity search
@@ -38,6 +44,7 @@ index.add(vectors)
 # Load the model for later use in embeddings
 model = sentence_transformers.SentenceTransformer(EMBEDDING_MODEL_NAME)
 
+
 # Define the search function
 def search(query: str, k: int) -> tuple[str]:
     query = numpy.expand_dims(model.encode(query), axis=0)
@@ -45,40 +52,66 @@ def search(query: str, k: int) -> tuple[str]:
     D, I = index.search(query, k)
     top_five = data.loc[I[0]]
 
-    search_results = "You are an AI assistant who delights in helping people" \
-        + "learn about research from the Design Research Collective. Here are" \
-        + "several abstracts from really cool, and really relevant, papers:\n\n"
+    search_results = (
+        "You are an AI assistant who delights in helping people learn about research from the Design "
+        "Research Collective. Here are several abstracts from really cool, and really relevant, "
+        "papers:\n\n"
+    )
 
     references = "\n\n## References\n\n"
 
     for i in range(k):
         search_results += top_five["bib_dict"].values[i]["abstract"] + "\n"
-        references += str(i+1) + ". " + ", ".join([author.split(" ")[-1] for author in top_five["bib_dict"].values[i]["author"].split(" and ")]) + ". (" + str(int(top_five["bib_dict"].values[i]["pub_year"])) +  "). [" + top_five["bib_dict"].values[i]["title"] + "]" \
-            + "(https://scholar.google.com/citations?view_op=view_citation&citation_for_view=" + top_five["author_pub_id"].values[i] + ").\n"
-
-    search_results += "\nIf these abstract aren't relevant to the followign query, please say that there is not much research in that area. Response to the following query from the perspective of the provided abstracts only:"
+        references += (
+            str(i + 1)
+            + ". "
+            + ", ".join(
+                [
+                    author.split(" ")[-1]
+                    for author in top_five["bib_dict"]
+                    .values[i]["author"]
+                    .split(" and ")
+                ]
+            )
+            + ". ("
+            + str(int(top_five["bib_dict"].values[i]["pub_year"]))
+            + "). ["
+            + top_five["bib_dict"].values[i]["title"]
+            + "]"
+            + "(https://scholar.google.com/citations?view_op=view_citation&citation_for_view="
+            + top_five["author_pub_id"].values[i]
+            + ").\n"
+        )
+
+    search_results += (
+        "\nIf these abstract aren't relevant to the following query, please reply 'I am unsure' or "
+        "similar. Respond to the following query from the perspective of the provided abstracts only:"
+    )
 
     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)
+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"
+    LLM_MODEL_NAME, torch_dtype="auto", device_map="auto"
 )
 
+
 def preprocess(message: str) -> tuple[str]:
     """Applies a preprocessing step to the user's message before the LLM receives it"""
     block_search_results, formatted_search_results = search(message, 5)
     return block_search_results + message, formatted_search_results
 
+
 def postprocess(response: str, bypass_from_preprocessing: str) -> str:
     """Applies a postprocessing step to the LLM's response before the user receives it"""
     return response + bypass_from_preprocessing
 
+
 @spaces.GPU
 def predict(message: str, history: list[str]) -> str:
     """This function is responsible for crafting a response"""
@@ -92,29 +125,23 @@ def predict(message: str, history: list[str]) -> str:
             history = history[-1]
     print(history)
     history_transformer_format = [
-        {"role": "assistant" if idx&1 else "user", "content": msg}
+        {"role": "assistant" if idx & 1 else "user", "content": msg}
         for idx, msg in enumerate(history)
     ] + [{"role": "user", "content": message}]
 
     # Stream a response from pipe
     text = tokenizer.apply_chat_template(
-        history_transformer_format,
-        tokenize=False,
-        add_generation_prompt=True
+        history_transformer_format, tokenize=False, add_generation_prompt=True
     )
     model_inputs = tokenizer([text], return_tensors="pt").to("cuda:0")
 
-    generate_kwargs = dict(
-        model_inputs,
-        streamer=streamer,
-        max_new_tokens=512
-    )
+    generate_kwargs = dict(model_inputs, streamer=streamer, max_new_tokens=512)
     t = threading.Thread(target=chatmodel.generate, kwargs=generate_kwargs)
     t.start()
 
     partial_message = ""
     for new_token in streamer:
-        if new_token != '<':
+        if new_token != "<":
             partial_message += new_token
             time.sleep(0.05)
             yield partial_message
@@ -124,21 +151,18 @@ def predict(message: str, history: list[str]) -> str:
 
 # Create and run the gradio interface
 gradio.ChatInterface(
-    predict, 
+    predict,
     examples=[
         "Tell me about new research at the intersection of additive manufacturing and machine learning",
         "What is a physics-informed neural network and what can it be used for?",
-        "What can agent-based models do about climate change?"
+        "What can agent-based models do about climate change?",
     ],
-    chatbot = gradio.Chatbot(    
-        show_label=False,
-        show_copy_button=True,
-        value=[["", GREETING]]
+    chatbot=gradio.Chatbot(
+        show_label=False, show_copy_button=True, value=[["", GREETING]]
     ),
-    retry_btn = None,
-    undo_btn = None,
-    clear_btn = None,
-    theme = "monochrome",
-    cache_examples = True,
-    fill_height = True,
-).launch(debug=True)
+    retry_btn=None,
+    undo_btn=None,
+    clear_btn=None,
+    cache_examples=True,
+    fill_height=True,
+).launch(debug=True)