Added token to space

core.py

@@ -1,6 +1,21 @@
 import streamlit as st
 from sentence_transformers import SentenceTransformer
+from huggingface_hub import HfApi, HfFolder
 import datasets
+import os
+
+
+@st.cache_data
+def login():
+    if not 'logged' in st.session_state:
+        st.session_state['logged'] = True
+        HF_TOKEN = os.environ.get("HF_TOKEN")
+        api=HfApi()
+        api.set_access_token(HF_TOKEN)
+        folder = HfFolder()
+        folder.save_token(HF_TOKEN)
+        return True
+    return False
 
 
 @st.cache_resource
@@ -10,7 +25,11 @@ def load_model():
 
 @st.cache_resource
 def load_index():
-    index = datasets.load_dataset("eremeev-d/arxiv-abstracts-small", use_auth_token=True)["train"]
+    index = datasets.load_dataset(
+        "eremeev-d/arxiv-abstracts-small", 
+        use_auth_token=True,
+        split="train"
+    )
     index.add_faiss_index('embedding')
     return index
 
@@ -29,7 +48,8 @@ def display_answer(query):
     for answer_id in range(len(answers)):
         with st.container():
             href = "https://arxiv.org/abs/{}".format(answers['id'][answer_id])
-            title = "<h3><a href=\"{}\">{}</a></h3>".format(href, answers['title'][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

@@ -1,3 +1,4 @@
 faiss-cpu~=1.7.2
 sentence-transformers~=2.2.2
-datasets~=2.10.1
+datasets~=2.10.1
+huggingface_hub~=0.10.1