Basic search on small dataset

core.py

@@ -1,11 +1,35 @@
 import streamlit as st
-import random
+from sentence_transformers import SentenceTransformer
+import datasets
+
+
+@st.cache_resource
+def load_model():
+    return SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+
+
+@st.cache_resource
+def load_index():
+    index = datasets.load_dataset("eremeev-d/arxiv-abstracts-small", use_auth_token=True)["train"]
+    index.add_faiss_index('embedding')
+    return index
+
+
+def get_answers(query):
+    model = load_model()
+    index = load_index()
+    query_embedding = model.encode(query)
+    scores, answers = index.get_nearest_examples('embedding', query_embedding)
+    return answers
 
 
 def display_answer(query):
     st.write("---")
-    for doc_id in range(random.randint(3, 5)):
+    answers = get_answers(query)
+    for answer_id in range(len(answers)):
         with st.container():
-            st.subheader("Paper #{}".format(doc_id))
-            st.markdown(r"Let's say it is description with LaTeX $\alpha$")
+            href = "https://arxiv.org/abs/{}".format(answers['id'][answer_id])
+            title = "<h3><a href=\"{}\">{}</a></h3>".format(href, answers['title'][answer_id])
+            st.write(title, unsafe_allow_html=True)
+            st.markdown(answers['abstract'][answer_id])
         st.write("---")

requirements.txt

@@ -0,0 +1 @@
+faiss-cpu~=1.7.2