Create main.py

main.py

@@ -0,0 +1,41 @@
+import gradio
+import numpy
+import pandas
+import sentence_transformers
+import datasets
+import faiss
+
+model = sentence_transformers.SentenceTransformer('allenai-specter')
+
+data = datasets.load_dataset("ccm/publications")['train'].to_pandas()
+
+dimensionality = len(data['embedding'][0])
+index = faiss.IndexFlatL2(dimensionality) 
+
+vectors = numpy.stack(data['embedding'].to_list(), axis=0)
+
+index.add(vectors)
+
+def search(query):
+    k=5
+    query = numpy.expand_dims(model.encode(query), axis=0)
+    _, I = top_five = 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]["citation"]
+        if top_five["pub_url"].values[i] is not None:
+            search_results += " [Full Paper](" + top_five["pub_url"].values[i] + ")"
+        search_results += "\n"
+    return search_results
+
+with gradio.Blocks() as demo:
+    query = gradio.Textbox(placeholder="Enter search terms...")
+    btn = gradio.Button("Search")
+    results = gradio.Markdown()
+    btn.click(fn=search, inputs=[query], outputs=results)
+
+demo.launch(debug=True)