Update main.py

main.py

@@ -1,17 +1,20 @@
+import json
+
 import gradio
+import datasets
+
 import numpy
 import pandas
 import sentence_transformers
-import datasets
 import faiss
 
 model = sentence_transformers.SentenceTransformer('allenai-specter')
 
 full_data = datasets.load_dataset("ccm/publications")['train'].to_pandas()
 
-substring = "0P9w_S0AAAAJ:yB1At4FlUx8C"
-filter = full_data['author_pub_id'].str.contains(substring)
-data = full_data[~filter]
+filter = ["\"abstract\": null" in json.dumps(bibdict) for bibdict in full_data['bibtex'].values]
+data = full_data[~pandas.Series(filter)]
+data.reset_index(inplace=True)
 
 dimensionality = len(data['embedding'][0])
 index = faiss.IndexFlatL2(dimensionality) 
@@ -22,19 +25,21 @@ index.add(vectors)
 
 def search(query, k):
     query = numpy.expand_dims(model.encode(query), axis=0)
-    _, I = top_five = index.search(query, k)
+    _, I = index.search(query, k)
     top_five = data.loc[I[0]]
     search_results = ""
 
     for i in range(k):
-        search_results += str(i+1) + ". "
-        search_results += '"' + top_five["bibtex"].values[i]["title"] + '" '
+        search_results += '## ' + top_five["bibtex"].values[i]["title"] + '\n'
         search_results += top_five["bibtex"].values[i]["citation"]
         if top_five["pub_url"].values[i] is not None:
-            search_results += " [Paper](" + top_five["pub_url"].values[i] + ")"
-        search_results += "\n"
+            search_results += ", [Paper](" + top_five["pub_url"].values[i] + ")"
+        search_results += "\t\n```\n"
+        search_results += json.dumps(top_five["bibtex"].values[i], indent=4)
+        search_results += "\t\n```\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)