trying to make things a bit faster

app.py

@@ -131,6 +131,10 @@ if 'ids' not in st.session_state:
     st.session_state.cites = arxiv_corpus['cites']
     st.session_state.years = arxiv_corpus['date']
     st.session_state.kws = arxiv_corpus['keywords']
+    st.session_state.ads_kws = arxiv_corpus['ads_keywords']
+    st.session_state.bibcode = arxiv_corpus['bibcode']
+    st.session_state.umap_x = arxiv_corpus['umap_x']
+    st.session_state.umap_y = arxiv_corpus['umap_y']
     st.toast('done caching. time taken: %.2f sec' %(time.time()-ts))
 
 
@@ -194,7 +198,7 @@ class EmbeddingRetrievalSystem():
         self.abstract = st.session_state.abstracts
         self.client = OpenAI(api_key = openai_key)
         self.embed_model = "text-embedding-3-small"
-        self.dataset = arxiv_corpus
+        self.dataset = st.session_state.arxiv_corpus
         self.kws = st.session_state.kws
         self.cites = st.session_state.cites
 
@@ -495,7 +499,7 @@ def Library(query, top_k = 7):
     rs = get_topk(query, top_k = top_k)
     op_docs = ''
     for paperno, i in enumerate(rs):
-        op_docs = op_docs + 'Paper %.0f:' %(paperno+1) +' (published in '+st.session_state.arxiv_corpus['bibcode'][i][0:4] + ') ' + st.session_state.titles[i]  + '\n' + st.session_state.abstracts[i] + '\n\n'
+        op_docs = op_docs + 'Paper %.0f:' %(paperno+1) +' (published in '+st.session_state.bibcode[i][0:4] + ') ' + st.session_state.titles[i]  + '\n' + st.session_state.abstracts[i] + '\n\n'
 
     return op_docs
 
@@ -504,7 +508,7 @@ def Library2(query, top_k = 7):
     absts, fnames = [], []
     for paperno, i in enumerate(rs):
         absts.append(st.session_state.abstracts[i])
-        fnames.append(st.session_state.arxiv_corpus['bibcode'][i])
+        fnames.append(st.session_state.bibcode[i])
     return absts, fnames, rs
 
 def get_paper_df(ids):
@@ -513,10 +517,10 @@ def get_paper_df(ids):
     for i in ids:
         papers.append(st.session_state.titles[i])
         scores.append(ids[i])
-        links.append('https://ui.adsabs.harvard.edu/abs/'+st.session_state.arxiv_corpus['bibcode'][i]+'/abstract')
-        yrs.append(st.session_state.arxiv_corpus['bibcode'][i][0:4])
-        cites.append(st.session_state.arxiv_corpus['cites'][i])
-        kws.append(st.session_state.arxiv_corpus['ads_keywords'][i])
+        links.append('https://ui.adsabs.harvard.edu/abs/'+st.session_state.bibcode[i]+'/abstract')
+        yrs.append(st.session_state.bibcode[i][0:4])
+        cites.append(st.session_state.cites[i])
+        kws.append(st.session_state.ads_kws[i])
 
     return pd.DataFrame({
         'Title': papers,
@@ -587,10 +591,10 @@ def create_embedding_plot(rs):
     """
 
     pltsource = ColumnDataSource(data=dict(
-        x=st.session_state.arxiv_corpus['umap_x'],
-        y=st.session_state.arxiv_corpus['umap_y'],
+        x=st.session_state.umap_x,
+        y=st.session_state.umap_y,
         title=st.session_state.titles,
-        link=st.session_state.arxiv_corpus['bibcode'],
+        link=st.session_state.bibcode,
     ))
 
     rsflag = np.zeros((len(st.session_state.ids),))
@@ -995,8 +999,8 @@ def main():
                            column_config = {'ADS Link':st.column_config.LinkColumn(display_text= 'https://ui.adsabs.harvard.edu/abs/(.*?)/abstract')}
                            )
 
-        with st.expander("Embedding map", expanded=False):
-            st.bokeh_chart(embedding_plot)
+        # with st.expander("Embedding map", expanded=False):
+        st.bokeh_chart(embedding_plot)
 
         col1, col2 = st.columns(2)