Update main.py

main.py

@@ -1,56 +1,78 @@
-import json
+import json  # For stringifying a dict
 
-import gradio
-import datasets
+import gradio  # GUI framework
+import datasets  # Used to load publication dataset
 
-import numpy
-import pandas
-import sentence_transformers
-import faiss
+import numpy  # For a few simple matrix operations
+import pandas  # Needed for operating on dataset
+import sentence_transformers  # Needed for query embedding
+import faiss  # Needed for fast similarity search
 
-model = sentence_transformers.SentenceTransformer('allenai-specter')
+# Load the model for later use in embeddings
+model = sentence_transformers.SentenceTransformer("allenai-specter")
 
-full_data = datasets.load_dataset("ccm/publications")['train'].to_pandas()
+# Load the dataset and convert to pandas
+full_data = datasets.load_dataset("ccm/publications")["train"].to_pandas()
 
-filter = ["\"abstract\": null" in json.dumps(bibdict) for bibdict in full_data['bib_dict'].values]
+# Define the base URL for Google Scholar
+SCHOLAR_URL = "https://scholar.google.com"
+
+# Filter out any publications without an abstract
+filter = [
+    '"abstract": null' in json.dumps(bibdict)
+    for bibdict in full_data["bib_dict"].values
+]
 data = full_data[~pandas.Series(filter)]
 data.reset_index(inplace=True)
 
-dimensionality = len(data['embedding'][0])
-index = faiss.IndexFlatL2(dimensionality) 
+# Create a FAISS index for fast similarity search
+index = faiss.IndexFlatL2(len(data["embedding"][0]))
 index.metric_type = faiss.METRIC_INNER_PRODUCT
+index.add(numpy.stack(data["embedding"].to_list(), axis=0))
 
-vectors = numpy.stack(data['embedding'].to_list(), axis=0)
-
-index.add(vectors)
 
-def search(query, k):
+# Define the search function
+def search(query: str, k: int):
     query = numpy.expand_dims(model.encode(query), axis=0)
     _, I = index.search(query, k)
     top_five = data.loc[I[0]]
     search_results = ""
 
     for i in range(k):
-        search_results += '### ' + top_five["bib_dict"].values[i]["title"] + '\n\n'
+        search_results += "### " + top_five["bib_dict"].values[i]["title"] + "\n\n"
         if top_five["pub_url"].values[i] is not None:
             search_results += "[Full Text](" + top_five["pub_url"].values[i] + ") "
         if top_five["citedby_url"].values[i] is not None:
-            search_results += "[Cited By](https://scholar.google.com" + top_five["citedby_url"].values[i] + ") "
+            search_results += (
+                "[Cited By](" + SCHOLAR_URL + top_five["citedby_url"].values[i] + ") "
+            )
         if top_five["url_related_articles"].values[i] is not None:
-            search_results += "[Related Articles](https://scholar.google.com" + top_five["url_related_articles"].values[i] + ") "
+            search_results += (
+                "[Related Articles]("
+                + SCHOLAR_URL
+                + top_five["url_related_articles"].values[i]
+                + ") "
+            )
         search_results += "\n\n```bibtex\n"
-        search_results += json.dumps(top_five["bibtex"].values[i], indent=4).replace('\\n', '\n').replace('\\t', '\t').strip("\"")
+        search_results += (
+            json.dumps(top_five["bibtex"].values[i], indent=4)
+            .replace("\\n", "\n")
+            .replace("\\t", "\t")
+            .strip('"')
+        )
         search_results += "```\n"
     return search_results
 
 
 with gradio.Blocks() as demo:
     with gradio.Group():
-        query = gradio.Textbox(placeholder="Enter search terms...", show_label=False, lines=1, max_lines=1)
+        query = gradio.Textbox(
+            placeholder="Enter search terms...", show_label=False, lines=1, max_lines=1
+        )
         with gradio.Accordion("Settings", open=False):
             k = gradio.Number(5.0, label="Number of results", precision=0)
     results = gradio.Markdown()
     query.change(fn=search, inputs=[query, k], outputs=results)
     k.change(fn=search, inputs=[query, k], outputs=results)
 
-demo.launch(debug=True)
+demo.launch(debug=True)