refactored code for speed

app.py

@@ -1,6 +1,9 @@
 import streamlit as st
 st.set_page_config(layout="wide")
 
+openai_key = st.secrets["openai_key"]
+cohere_key = st.secrets['cohere_key']
+
 import numpy as np
 from abc import ABC, abstractmethod
 from typing import List, Dict, Any, Tuple
@@ -28,6 +31,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_core.output_parsers import StrOutputParser
 from langchain.callbacks import FileCallbackHandler
 from langchain.callbacks.manager import CallbackManager
+from langchain.schema import Document
 
 import instructor
 from pydantic import BaseModel, Field
@@ -43,121 +47,51 @@ import faiss
 import spacy
 from string import punctuation
 import pytextrank
-
-from bokeh.plotting import figure
-from bokeh.models import ColumnDataSource
-from bokeh.io import output_notebook
-from bokeh.palettes import Spectral5
-from bokeh.transform import linear_cmap
+from prompts import *
 
 ts = time.time()
-st.session_state.ts = ts
-
-openai_key = st.secrets["openai_key"]
-cohere_key = st.secrets['cohere_key']
 
-if 'nlp' not in st.session_state:
+@st.cache_resource
+def load_nlp():
     nlp = spacy.load("en_core_web_sm")
     nlp.add_pipe("textrank")
-    st.session_state.nlp = nlp
-
-try:
-    stopwords.words('english')
-except:
-    nltk.download('stopwords')
-    stopwords.words('english')
+    try:
+        stopwords.words('english')
+    except:
+        nltk.download('stopwords')
+        stopwords.words('english')
+    return nlp
+
+# @st.cache_resource
+# def load_embeddings():
+#     return OpenAIEmbeddings(model="text-embedding-3-small", api_key=st.secrets["openai_key"])
+#
+# @st.cache_resource
+# def load_llm():
+#     return ChatOpenAI(temperature=0, model_name='gpt-4o-mini', openai_api_key=st.secrets["openai_key"])
 
 st.session_state.gen_llm = openai_llm(temperature=0,
-                     model_name='gpt-4o-mini', 
+                     model_name='gpt-4o-mini',
                      openai_api_key = openai_key)
 st.session_state.consensus_client = instructor.patch(OpenAI(api_key=openai_key))
 st.session_state.embed_client = OpenAI(api_key = openai_key)
 embed_model = "text-embedding-3-small"
 st.session_state.embeddings = OpenAIEmbeddings(model = embed_model, api_key = openai_key)
 
-st.image('local_files/pathfinder_logo.png')
-
-st.expander("What is Pathfinder / How do I use it?", expanded=False).write(
-        """
-        Pathfinder v2.0 is a framework for searching and visualizing astronomy papers on the [arXiv](https://arxiv.org/) and [ADS](https://ui.adsabs.harvard.edu/) using the context
-        sensitivity from modern large language models (LLMs) to better parse patterns in paper contexts.
-
-        This tool was built during the [JSALT workshop](https://www.clsp.jhu.edu/2024-jelinek-summer-workshop-on-speech-and-language-technology/) to do awesome things.
-
-        **👈 Use the sidebar to tweak the search parameters to get better results**.
-
-        ### Tool summary:
-        - Please wait while the initial data loads and compiles, this takes about a minute initially.
-
-        This is not meant to be a replacement to existing tools like the
-        [ADS](https://ui.adsabs.harvard.edu/),
-        [arxivsorter](https://www.arxivsorter.org/), semantic search or google scholar, but rather a supplement to find papers
-        that otherwise might be missed during a literature survey.
-        It is trained on astro-ph (astrophysics of galaxies) papers up to last-year-ish mined from arxiv and supplemented with ADS metadata,
-        if you are interested in extending it please reach out!
-
-        Also add: feedback form, socials, literature, contact us, copyright, collaboration, etc.
-
-        The image below shows a representation of all the astro-ph.GA papers that can be explored in more detail
-        using the `Arxiv embedding` page. The papers tend to cluster together by similarity, and result in an
-        atlas that shows well studied (forests) and currently uncharted areas (water).
-        """
-    )
-
-
-st.sidebar.header("Fine-tune the search")
-top_k = st.sidebar.slider("Number of papers to retrieve:", 3, 30, 10)
-extra_keywords = st.sidebar.text_input("Enter extra keywords (comma-separated):")
-
-st.sidebar.subheader("Toggles")
-toggle_a = st.sidebar.toggle("Weight by keywords", value = False)
-toggle_b = st.sidebar.toggle("Weight by date", value = False)
-toggle_c = st.sidebar.toggle("Weight by citations", value = False)
-
-method = st.sidebar.radio("Retrieval method:", ["Semantic search", "Semantic search + HyDE", "Semantic search + HyDE + CoHERE"], index=2)
-
-method2 = st.sidebar.radio("Generation complexity:", ["Basic RAG","ReAct Agent"])
-
-question_type = st.sidebar.selectbox("Select question type:", ["Multi-paper (Default)", "Single-paper", "Bibliometric", "Broad but nuanced"])
-st.session_state.question_type = question_type
-# store_output = st.sidebar.button("Save output")
-
-query = st.text_input("Ask me anything:")
-submit_button = st.button("Run pathfinder!")
-
-search_text_list = ['rooting around in the paper pile...','looking for clarity...','scanning the event horizon...','peering into the abyss...','potatoes power this ongoing search...']
 
-if 'arxiv_corpus' not in st.session_state:
-    with st.spinner('loading data (please wait for this to finish before querying)...'):
-        # try:
+# @st.cache_data
+def load_arxiv_corpus():
+    with st.spinner('loading astro-ph corpus'):
         arxiv_corpus = load_from_disk('data/')
-        # except:
-        #     st.write('downloading data')
-        #     arxiv_corpus = load_dataset('kiyer/pathfinder_arxiv_data',split='train')
-        #     # arxiv_corpus = load_dataset('kiyer/pathfinder_arxiv_data_galaxy',split='train')
-        #     arxiv_corpus.save_to_disk('data/')
-        arxiv_corpus.add_faiss_index('embed')
-        st.session_state.arxiv_corpus = arxiv_corpus
-        st.toast('loaded arxiv corpus')
-
-        if 'ids' not in st.session_state:
-            with st.spinner('making the LLM talk to the astro papers...'):
-                st.session_state.ids = st.session_state.arxiv_corpus['ads_id']
-                st.session_state.titles = st.session_state.arxiv_corpus['title']
-                st.session_state.abstracts = st.session_state.arxiv_corpus['abstract']
-                st.session_state.authors = st.session_state.arxiv_corpus['authors']
-                st.session_state.cites = st.session_state.arxiv_corpus['cites']
-                st.session_state.years = st.session_state.arxiv_corpus['date']
-                st.session_state.kws = st.session_state.arxiv_corpus['keywords']
-                st.session_state.ads_kws = st.session_state.arxiv_corpus['ads_keywords']
-                st.session_state.bibcode = st.session_state.arxiv_corpus['bibcode']
-                st.session_state.umap_x = st.session_state.arxiv_corpus['umap_x']
-                st.session_state.umap_y = st.session_state.arxiv_corpus['umap_y']
-                st.toast('done caching. time taken: %.2f sec' %(time.time()-ts))   
+        arxiv_corpus.load_faiss_index('embed', 'data/astrophindex.faiss')
+        st.toast('loaded data. time taken: %.2f sec' %(time.time()-ts))
+    return arxiv_corpus
 
 def get_keywords(text):
     result = []
     pos_tag = ['PROPN', 'ADJ', 'NOUN']
+    if 'nlp' not in st.session_state:
+        st.session_state.nlp = load_nlp()
     doc = st.session_state.nlp(text.lower())
     for token in doc:
         if(token.text in st.session_state.nlp.Defaults.stop_words or token.text in punctuation):
@@ -166,43 +100,19 @@ def get_keywords(text):
             result.append(token.text)
     return result
 
-def parse_doc(text, nret = 10):
-    local_kws = []
-    doc = st.session_state.nlp(text)
-    # examine the top-ranked phrases in the document
-    for phrase in doc._.phrases[:nret]:
-        # print(phrase.text)
-        local_kws.append(phrase.text)
-    return local_kws
 
-class EmbeddingRetrievalSystem():
 
-    def __init__(self, weight_citation = False, weight_date = False, weight_keywords = False):
 
-        self.ids = st.session_state.ids
-        self.years = st.session_state.years
-        self.abstract = st.session_state.abstracts
+class RetrievalSystem():
+
+    def __init__(self):
+
+        self.dataset = st.session_state.arxiv_corpus
         self.client = OpenAI(api_key = openai_key)
         self.embed_model = "text-embedding-3-small"
-        self.dataset = st.session_state.arxiv_corpus
-        self.kws = st.session_state.kws
-        self.cites = st.session_state.cites
-
-        self.weight_citation = weight_citation
-        self.weight_date = weight_date
-        self.weight_keywords = weight_keywords
-        self.id_to_index = {self.ids[i]: i for i in range(len(self.ids))}
-        
         self.generation_client = openai_llm(temperature=0,model_name='gpt-4o-mini', openai_api_key = openai_key)
-
-        # self.citation_filter = CitationFilter(self.dataset)
-        # self.date_filter = DateFilter(self.dataset['date'])
-        # self.keyword_filter = KeywordFilter(corpus=self.dataset, remove_capitals=True)
-
-    def parse_date(self, id):
-        # indexval = np.where(self.ids == id)[0][0]
-        indexval = id
-        return self.years[indexval]
+        self.hyde_client = openai_llm(temperature=0.5,model_name='gpt-4o-mini', openai_api_key = openai_key)
+        self.cohere_client = cohere.Client(cohere_key)
 
     def make_embedding(self, text):
         str_embed = self.client.embeddings.create(input = [text], model = self.embed_model).data[0].embedding
@@ -215,39 +125,29 @@ class EmbeddingRetrievalSystem():
     def get_query_embedding(self, query):
         return self.make_embedding(query)
 
-    def analyze_temporal_query(self, query):
-        return
-
     def calc_faiss(self, query_embedding, top_k = 100):
         # xq = query_embedding.reshape(-1,1).T.astype('float32')
         # D, I = self.index.search(xq, top_k)
         # return I[0], D[0]
         tmp = self.dataset.search('embed', query_embedding, k=top_k)
-        return [tmp.indices, tmp.scores]
-
-    def rank_and_filter(self, query, query_embedding, query_date, top_k = 10, return_scores=False, time_result=None):
-
-        # st.write('status')
+        return [tmp.indices, tmp.scores, self.dataset[tmp.indices]]
 
-        # st.write('toggles', self.toggles)
-        # st.write('question_type', self.question_type)
-        # st.write('rag method', self.rag_method)
-        # st.write('gen method', self.gen_method)
+    def rank_and_filter(self, query, query_embedding, top_k = 10, top_k_internal = 1000, return_scores=False):
 
         self.weight_keywords = self.toggles["Keyword weighting"]
         self.weight_date = self.toggles["Time weighting"]
         self.weight_citation = self.toggles["Citation weighting"]
 
-        topk_indices, similarities = self.calc_faiss(np.array(query_embedding), top_k = 1000)
+        topk_indices, similarities, small_corpus = self.calc_faiss(np.array(query_embedding), top_k = top_k_internal)
         similarities = 1/similarities # converting from a distance (less is better) to a similarity (more is better)
 
-        query_kws = get_keywords(query)
-        input_kws = self.query_input_keywords
-        query_kws = query_kws + input_kws
-        self.query_kws = query_kws
-
         if self.weight_keywords == True:
-            sub_kws = [self.kws[i] for i in topk_indices]
+
+            query_kws = get_keywords(query)
+            input_kws = self.query_input_keywords
+            query_kws = query_kws + input_kws
+            self.query_kws = query_kws
+            sub_kws = [small_corpus['keywords'][i] for i in range(top_k_internal)]
             kw_weight = np.zeros((len(topk_indices),)) + 0.1
 
             for k in query_kws:
@@ -263,7 +163,7 @@ class EmbeddingRetrievalSystem():
             kw_weight = np.ones((len(topk_indices),))
 
         if self.weight_date == True:
-            sub_dates = [self.years[i] for i in topk_indices]
+            sub_dates = [small_corpus['date'][i] for i in range(top_k_internal)]
             date = datetime.now().date()
             date_diff = np.array([((date - i).days / 365.) for i in sub_dates])
             # age_weight = (1 + np.exp(date_diff/2.1))**(-1) + 0.5
@@ -274,7 +174,7 @@ class EmbeddingRetrievalSystem():
 
         if self.weight_citation == True:
             # st.write('weighting by citations')
-            sub_cites = np.array([self.cites[i] for i in topk_indices])
+            sub_cites = np.array([small_corpus['cites'][i] for i in range(top_k_internal)])
             temp = sub_cites.copy()
             temp[sub_cites > 300] = 300.
             cite_weight = (1 + np.exp((300-temp)/42.0))**(-1.)
@@ -287,102 +187,24 @@ class EmbeddingRetrievalSystem():
         filtered_results = [[topk_indices[i], similarities[i]] for i in range(len(similarities))]
         top_results = sorted(filtered_results, key=lambda x: x[1], reverse=True)[:top_k]
 
+        top_scores = [doc[1] for doc in top_results]
+        top_indices = [doc[0] for doc in top_results]
+        small_df = self.dataset[top_indices]
+
         if return_scores:
-            return {doc[0]: doc[1] for doc in top_results}
+            return {doc[0]: doc[1] for doc in top_results}, small_df
 
         # Only keep the document IDs
         top_results = [doc[0] for doc in top_results]
-        return top_results
-
-    def retrieve(self, query, top_k, time_result=None, query_date = None, return_scores = False):
-
-        query_embedding = self.get_query_embedding(query)
-
-        # Judge time relevance
-        if time_result is None:
-            if self.weight_date:
-                time_result, time_taken = self.analyze_temporal_query(query, self.anthropic_client)
-            else:
-                time_result = {'has_temporal_aspect': False, 'expected_year_filter': None, 'expected_recency_weight': None}
-
-        top_results = self.rank_and_filter(query,
-                                           query_embedding,
-                                           query_date,
-                                           top_k,
-                                           return_scores = return_scores,
-                                           time_result = time_result)
-
-        return top_results
-
-class HydeRetrievalSystem(EmbeddingRetrievalSystem):
-    def __init__(self, generation_model: str = "claude-3-haiku-20240307",
-                 embedding_model: str = "text-embedding-3-small",
-             temperature: float = 0.5,
-                 max_doclen: int = 500,
-                 generate_n: int = 1,
-                 embed_query = True,
-                 conclusion = False, **kwargs):
-
-        # Handle the kwargs for the superclass init -- filters/citation weighting
-        super().__init__(**kwargs)
-
-        if max_doclen * generate_n > 8191:
-            raise ValueError("Too many tokens. Please reduce max_doclen or generate_n.")
-
-        self.embedding_model = embedding_model
-        self.generation_model = generation_model
-
-        # HYPERPARAMETERS
-        self.temperature = temperature # generation temperature
-        self.max_doclen = max_doclen # max tokens for generation
-        self.generate_n = generate_n # how many documents
-        self.embed_query = embed_query # embed the query vector?
-        self.conclusion = conclusion # generate conclusion as well?
-
-        # self.anthropic_key = anthropic_key
-        # self.generation_client = anthropic.Anthropic(api_key = self.anthropic_key)
-        
-
-    def retrieve(self, query: str, top_k: int = 10, return_scores = False, time_result = None) -> List[Tuple[str, str, float]]:
-        if time_result is None:
-            if self.weight_date: time_result, time_taken = analyze_temporal_query(query, self.anthropic_client)
-            else: time_result = {'has_temporal_aspect': False, 'expected_year_filter': None, 'expected_recency_weight': None}
-
-        docs = self.generate_docs(query)
-        st.expander('Abstract generated with hyde', expanded=False).write(docs)
-
-        doc_embeddings = self.embed_docs(docs)
-
-        if self.embed_query:
-            query_emb = self.embed_docs([query])[0]
-            doc_embeddings.append(query_emb)
-
-        embedding = np.mean(np.array(doc_embeddings), axis = 0)
-
-        top_results = self.rank_and_filter(query, embedding, query_date=None, top_k = top_k, return_scores = return_scores, time_result = time_result)
-
-        return top_results
+        return top_results, small_df
 
     def generate_doc(self, query: str):
         prompt = """You are an expert astronomer. Given a scientific query, generate the abstract of an expert-level research paper
                             that answers the question. Stick to a maximum length of {} tokens and return just the text of the abstract and conclusion.
                             Do not include labels for any section. Use research-specific jargon.""".format(self.max_doclen)
-        # st.write('invoking hyde generation')
-
-        # message = self.generation_client.messages.create(
-        #         model = self.generation_model,
-        #         max_tokens = self.max_doclen,
-        #         temperature = self.temperature,
-        #         system = prompt,
-        #         messages=[{ "role": "user",
-        #                 "content": [{"type": "text", "text": query,}] }]
-        #     )
-        # return message.content[0].text
 
         messages = [("system",prompt,),("human", query),]
-        return self.generation_client.invoke(messages).content
-
-
+        return self.hyde_client.invoke(messages).content
 
     def generate_docs(self, query: str):
         docs = []
@@ -393,120 +215,185 @@ class HydeRetrievalSystem(EmbeddingRetrievalSystem):
     def embed_docs(self, docs: List[str]):
         return self.embed_batch(docs)
 
-class HydeCohereRetrievalSystem(HydeRetrievalSystem):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
+    def retrieve(self, query, top_k, return_scores = False,
+                 embed_query=True, max_doclen=250,
+                 generate_n=1, temperature=0.5,
+                 rerank_top_k = 250):
 
-        self.cohere_key = cohere_key
-        self.cohere_client = cohere.Client(self.cohere_key)
+        if max_doclen * generate_n > 8191:
+            raise ValueError("Too many tokens. Please reduce max_doclen or generate_n.")
 
-    def retrieve(self, query: str,
-                 top_k: int = 10,
-                 rerank_top_k: int = 250,
-                 return_scores = False, time_result = None,
-                 reweight = False) -> List[Tuple[str, str, float]]:
+        query_embedding = self.get_query_embedding(query)
 
-        if time_result is None:
-            if self.weight_date: time_result, time_taken = analyze_temporal_query(query, self.anthropic_client)
-            else: time_result = {'has_temporal_aspect': False, 'expected_year_filter': None, 'expected_recency_weight': None}
+        if self.hyde == True:
+            self.max_doclen = max_doclen
+            self.generate_n = generate_n
+            self.hyde_client.temperature = temperature
+            self.embed_query = embed_query
+            docs = self.generate_docs(query)
+            st.expander('Abstract generated with hyde', expanded=False).write(docs)
+            doc_embeddings = self.embed_docs(docs)
+            if self.embed_query:
+                query_emb = self.embed_docs([query])[0]
+                doc_embeddings.append(query_emb)
+            query_embedding = np.mean(np.array(doc_embeddings), axis = 0)
+
+        if self.rerank == True:
+            top_results, small_df = self.rank_and_filter(query,
+                                           query_embedding,
+                                           rerank_top_k,
+                                           return_scores = False)
+            try:
+                docs_for_rerank = [small_df['abstract'][i] for i in range(rerank_top_k)]
+                if len(docs_for_rerank) == 0:
+                    return []
+                reranked_results = self.cohere_client.rerank(
+                    query=query,
+                    documents=docs_for_rerank,
+                    model='rerank-english-v3.0',
+                    top_n=top_k
+                )
+                final_results = []
+                for result in reranked_results.results:
+                    doc_id = top_results[result.index]
+                    doc_text = docs_for_rerank[result.index]
+                    score = float(result.relevance_score)
+                    final_results.append([doc_id, "", score])
+                final_indices = [doc[0] for doc in final_results]
+                if return_scores:
+                    return {result[0]: result[2] for result in final_results}, self.dataset[final_indices]
+                return [doc[0] for doc in final_results], self.dataset[final_indices]
+            except:
+                print('heavy load, please wait 10s and try again.')
+        else:
+            top_results, small_df = self.rank_and_filter(query,
+                                               query_embedding,
+                                               top_k,
+                                               return_scores = return_scores)
 
-        top_results = super().retrieve(query, top_k = rerank_top_k, time_result = time_result)
+        return top_results, small_df
 
-        # doc_texts = self.get_document_texts(top_results)
-        # docs_for_rerank = [f"Abstract: {doc['abstract']}\nConclusions: {doc['conclusions']}" for doc in doc_texts]
-        docs_for_rerank = [self.abstract[i] for i in top_results]
+    def return_formatted_df(self, top_results, small_df):
 
-        if len(docs_for_rerank) == 0:
-            return []
+        df = pd.DataFrame(small_df)
+        df = df.drop(columns=['embed','umap_x','umap_y','cite_bibcodes','ref_bibcodes'])
+        links = ['https://ui.adsabs.harvard.edu/abs/'+i+'/abstract' for i in small_df['bibcode']]
+        scores = [top_results[i] for i in top_results]
+        df.insert(1,'ADS Link',links,True)
+        df.insert(2,'Relevance',scores,True)
+        df = df[['ADS Link','Relevance','date','cites','title','authors','abstract','keywords','ads_id']]
+        return df
 
-        reranked_results = self.cohere_client.rerank(
-            query=query,
-            documents=docs_for_rerank,
-            model='rerank-english-v3.0',
-            top_n=top_k
-        )
+# @st.cache_resource
+def load_ret_system():
+    with st.spinner('loading retrieval system...'):
+        ec = RetrievalSystem()
+        st.toast('loaded retrieval system. time taken: %.2f sec' %(time.time()-ts))
+    return ec
+
+
+st.image('local_files/pathfinder_logo.png')
 
-        final_results = []
-        for result in reranked_results.results:
-            doc_id = top_results[result.index]
-            doc_text = docs_for_rerank[result.index]
-            score = float(result.relevance_score)
-            final_results.append([doc_id, "", score])
+st.expander("What is Pathfinder / How do I use it?", expanded=False).write(
+        """
+        Pathfinder v2.0 is a framework for searching and visualizing astronomy papers on the [arXiv](https://arxiv.org/) and [ADS](https://ui.adsabs.harvard.edu/) using the context
+        sensitivity from modern large language models (LLMs) to better parse patterns in paper contexts.
 
-        if reweight:
-            if time_result['has_temporal_aspect']:
-                final_results = self.date_filter.filter(final_results, time_score = time_result['expected_recency_weight'])
+        This tool was built during the [JSALT workshop](https://www.clsp.jhu.edu/2024-jelinek-summer-workshop-on-speech-and-language-technology/) to do awesome things.
 
-            if self.weight_citation: self.citation_filter.filter(final_results)
+        **👈 Use the sidebar to tweak the search parameters to get better results**.
 
-        if return_scores:
-            return {result[0]: result[2] for result in final_results}
+        ### Tool summary:
+        - Please wait while the initial data loads and compiles, this takes about a minute initially.
 
-        return [doc[0] for doc in final_results]
+        This is not meant to be a replacement to existing tools like the
+        [ADS](https://ui.adsabs.harvard.edu/),
+        [arxivsorter](https://www.arxivsorter.org/), semantic search or google scholar, but rather a supplement to find papers
+        that otherwise might be missed during a literature survey.
+        It is trained on astro-ph (astrophysics of galaxies) papers up to last-year-ish mined from arxiv and supplemented with ADS metadata,
+        if you are interested in extending it please reach out!
 
-    def embed_docs(self, docs: List[str]):
-        return self.embed_batch(docs)
+        Also add: feedback form, socials, literature, contact us, copyright, collaboration, etc.
+
+        The image below shows a representation of all the astro-ph.GA papers that can be explored in more detail
+        using the `Arxiv embedding` page. The papers tend to cluster together by similarity, and result in an
+        atlas that shows well studied (forests) and currently uncharted areas (water).
+        """
+    )
+
+st.sidebar.header("Fine-tune the search")
+top_k = st.sidebar.slider("Number of papers to retrieve:", 1, 30, 10)
+extra_keywords = st.sidebar.text_input("Enter extra keywords (comma-separated):")
+keywords = [kw.strip() for kw in extra_keywords.split(',')] if extra_keywords else []
 
-# --------- other fns ------------------
+st.sidebar.subheader("Toggles")
+toggle_a = st.sidebar.toggle("Weight by keywords", value = False)
+toggle_b = st.sidebar.toggle("Weight by date", value = False)
+toggle_c = st.sidebar.toggle("Weight by citations", value = False)
+toggles = {'Keyword weighting': toggle_a, 'Time weighting': toggle_b, 'Citation weighting': toggle_c}
+
+method = st.sidebar.radio("Retrieval method:", ["Semantic search", "Semantic search + HyDE", "Semantic search + HyDE + CoHERE"], index=2)
+method2 = st.sidebar.radio("Generation complexity:", ["Basic RAG","ReAct Agent"])
 
-def get_topk(query, top_k):
-    print('running retrieval')
-    rs = st.session_state.ec.retrieve(query, top_k, return_scores=True)
-    return rs
+st.session_state.top_k = top_k
+st.session_state.keywords = keywords
+st.session_state.toggles = toggles
+st.session_state.method = method
+st.session_state.method2 = method2
+
+if (method == "Semantic search"):
+    st.session_state.hyde = False
+    st.session_state.cohere = False
+elif (method == "Semantic search + HyDE"):
+    st.session_state.hyde = True
+    st.session_state.cohere = False
+elif (method == "Semantic search + HyDE + CoHERE"):
+    st.session_state.hyde = True
+    st.session_state.cohere = True
+
+if method2 == "Basic RAG":
+    st.session_state.gen_method = 'rag'
+elif method2 == "ReAct Agent":
+    st.session_state.gen_method = 'agent'
+
+question_type = st.sidebar.selectbox("Prompt specialization:", ["Multi-paper (Default)", "Single-paper", "Bibliometric", "Broad but nuanced"])
+st.session_state.question_type = question_type
+# store_output = st.sidebar.button("Save output")
+
+query = st.text_input("Ask me anything:")
+st.session_state.query = query
+st.write(query)
+submit_button = st.button("Run pathfinder!", key='runpfdr')
 
-def Library(query, top_k = 7):
-    rs = get_topk(query, top_k = top_k)
+search_text_list = ['rooting around in the paper pile...','looking for clarity...','scanning the event horizon...','peering into the abyss...','potatoes power this ongoing search...']
+gen_text_list = ['making the LLM talk to the papers...','invoking arcane rituals...','gone to library, please wait...','is there really an answer to this...']
+
+if 'arxiv_corpus' not in st.session_state:
+    st.session_state.arxiv_corpus = load_arxiv_corpus()
+
+# @st.fragment()
+def run_query_ret(query):
+    tr = time.time()
+    ec = load_ret_system()
+    ec.query_input_keywords = st.session_state.keywords
+    ec.toggles = st.session_state.toggles
+    ec.hyde = st.session_state.hyde
+    ec.rerank = st.session_state.cohere
+    rs, small_df = ec.retrieve(query, top_k = st.session_state.top_k, return_scores=True)
+    formatted_df = ec.return_formatted_df(rs, small_df)
+    st.toast('got top-k papers. time taken: %.2f sec' %(time.time()-tr))
+    return formatted_df
+
+def Library(query):
+    papers_df = run_query_ret(st.session_state.query)
     op_docs = ''
-    for paperno, i in enumerate(rs):
-        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'
+    for i in range(len(papers_df)):
+        op_docs = op_docs + 'Paper %.0f:' %(i+1) + papers_df['title'][i]  + '\n' + papers_df['abstract'][i] + '\n\n'
 
     return op_docs
 
-def Library2(query, top_k = 7):
-    rs = get_topk(query, top_k = top_k)
-    absts, fnames = [], []
-    for paperno, i in enumerate(rs):
-        absts.append(st.session_state.abstracts[i])
-        fnames.append(st.session_state.bibcode[i])
-    return absts, fnames, rs
-
-def get_paper_df(ids):
-
-    papers, scores, yrs, links, cites, kws, authors, absts = [], [], [], [], [], [], [], []
-    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.bibcode[i]+'/abstract')
-        yrs.append(st.session_state.bibcode[i][0:4])
-        cites.append(st.session_state.cites[i])
-        authors.append(st.session_state.authors[i][0])
-        kws.append(st.session_state.ads_kws[i])
-        absts.append(st.session_state.abstracts[i])
-
-    return pd.DataFrame({
-        'Title': papers,
-        'Relevance': scores,
-        'Lead author': authors,
-        'Year': yrs,
-        'ADS Link': links,
-        'Citations': cites,
-        'Keywords': kws,
-        'Abstract': absts
-    })
-
-def extract_keywords(question, ec):
-    # Simulated keyword extraction (replace with actual logic)
-    return ['keyword1', 'keyword2', 'keyword3']
-
-# Function to estimate consensus (replace with actual implementation)
-def estimate_consensus():
-    # Simulated consensus estimation (replace with actual calculation)
-    return 0.75
-
-
-def run_agent_qa(query, top_k):
-
-    # define tools
+def run_agent_qa(query):
+
     search = DuckDuckGoSearchAPIWrapper()
     tools = [
         Tool(
@@ -524,42 +411,8 @@ def run_agent_qa(query, top_k):
     if 'tools' not in st.session_state:
         st.session_state.tools = tools
 
-    # define prompt
-
-    # for another question type:
-    # First, find the quotes from the document that are most relevant to answering the question, and then print them in numbered order.
-    # Quotes should be relatively short. If there are no relevant quotes, write “No relevant quotes” instead.
-
-
-    template = """You are an expert astronomer and cosmologist.
-    Answer the following question as best you can using information from the library, but speaking in a concise and factual manner.
-    If you can not come up with an answer, say you do not know.
-    Try to break the question down into smaller steps and solve it in a logical manner.
-
-    You have access to the following tools:
-
-    {tools}
-
-    Use the following format:
-
-    Question: the input question you must answer
-    Thought: you should always think about what to do
-    Action: the action to take, should be one of [{tool_names}]
-    Action Input: the input to the action
-    Observation: the result of the action
-    ... (this Thought/Action/Action Input/Observation can repeat N times)
-    Thought: I now know the final answer
-    Final Answer: the final answer to the original input question. provide information about how you arrived at the answer, and any nuances or uncertainties the reader should be aware of
-
-    Begin! Remember to speak in a pedagogical and factual manner."
-
-    Question: {input}
-    Thought:{agent_scratchpad}"""
-
     prompt = hub.pull("hwchase17/react")
-    prompt.template=template
-
-    # path to write intermediate trace to
+    prompt.template = react_prompt
 
     file_path = "agent_trace.txt"
     try:
@@ -569,8 +422,6 @@ def run_agent_qa(query, top_k):
     file_handler = FileCallbackHandler(file_path)
     callback_manager=CallbackManager([file_handler])
 
-    # define and execute agent
-
     tool_names = [tool.name for tool in st.session_state.tools]
     if 'agent' not in st.session_state:
         # agent = ZeroShotAgent(llm_chain=llm_chain, allowed_tools=tool_names)
@@ -584,125 +435,27 @@ def run_agent_qa(query, top_k):
     answer = st.session_state.agent_executor.invoke({"input": query,})
     return answer
 
-regular_prompt = """You are an expert astronomer and cosmologist.
-Answer the following question as best you can using information from the library, but speaking in a concise and factual manner.
-If you can not come up with an answer, say you do not know.
-Try to break the question down into smaller steps and solve it in a logical manner.
+def run_rag_qa(query, papers_df):
 
-Provide information about how you arrived at the answer, and any nuances or uncertainties the reader should be aware of.
-
-Begin! Remember to speak in a pedagogical and factual manner."
-
-Relevant documents:{context}
-
-Question: {question}
-Answer:"""
-
-bibliometric_prompt = """You are an AI assistant with expertise in astronomy and astrophysics literature. Your task is to assist with relevant bibliometric information in response to a user question. The user question may consist of identifying key papers, authors, or trends in a specific area of astronomical research.
-
-Depending on what the user wants, direct them to consult the NASA Astrophysics Data System (ADS) at https://ui.adsabs.harvard.edu/. Provide them with the recommended ADS query depending on their question.
-
-Here's a more detailed guide on how to use NASA ADS for various types of queries:
-
-Basic topic search: Enter keywords in the search bar, e.g., "exoplanets". Use quotation marks for exact phrases, e.g., "dark energy”
-Author search: Use the syntax author:"Last Name, First Name", e.g., author:"Hawking, S". For papers by multiple authors, use AND, e.g., author:"Hawking, S" AND author:"Ellis, G"
-Date range: Use year:YYYY-YYYY, e.g., year:2010-2020. For papers since a certain year, use year:YYYY-, e.g., year:2015-
-4.Combining search terms: Use AND, OR, NOT operators, e.g., "black holes" AND (author:"Hawking, S" OR author:"Penrose, R")
-Filtering results: Use the left sidebar to filter by publication year, article type, or astronomy database
-Sorting results: Use the "Sort" dropdown menu to order by options like citation count, publication date, or relevance
-Advanced searches: Click on the "Search" dropdown menu and select "Classic Form" for field-specific searchesUse bibcode:YYYY for a specific journal/year, e.g., bibcode:2020ApJ to find all Astrophysical Journal papers from 2020
-Finding review articles: Wrap the query in the reviews() operator (e.g. reviews(“dark energy”)) 
-Excluding preprints: Add NOT doctype:"eprint" to your search
-Citation metrics: Click on the citation count of a paper to see its citation history and who has cited it
-
-Some examples:
-
-Example 1:
-“How many papers published in 2022 used data from MAST missions?”
-Your response should be: year:2022  data:"MAST"
-
-Example 2:
-“What are the most cited papers on spiral galaxy halos measured in X-rays, with publication date from 2010 to 2023?
-Your response should be: "spiral galaxy halos" AND "x-ray" year:2010-2024
-
-Example 3:
-“Can you list 3 papers published by “< name>” as first author?”
-Your response should be: author: “^X”
-
-Example 4:
-“Based on papers with “<name>” as an author or co-author, can you suggest the five most recent astro-ph papers that would be relevant?”
-Your response should be: 
-
-Remember to advise users that while these examples cover many common scenarios, NASA ADS has many more advanced features that can be explored through its documentation.
-
-Relevant documents:{context}
-Question: {question}
-
-Response:"""
-
-single_paper_prompt = """You are an astronomer with access to a vast database of astronomical facts and figures. Your task is to provide a concise, accurate answer to the following specific factual question about astronomy or astrophysics.
-Provide the requested information clearly and directly. If relevant, include the source of your information or any recent updates to this fact. If there's any uncertainty or variation in the accepted value, briefly explain why.
-If the question can't be answered with a single fact, provide a short, focused explanation. Always prioritize accuracy over speculation.
-Relevant documents:{context}
-Question: {question}
-Response:"""
-
-deep_knowledge_prompt = """You are an expert astronomer with deep knowledge across various subfields of astronomy and astrophysics. Your task is to provide a comprehensive and nuanced answer to the following question, which involves an unresolved topic or requires broad, common-sense understanding.
-Consider multiple perspectives and current debates in the field. Explain any uncertainties or ongoing research. If relevant, mention how this topic connects to other areas of astronomy.
-Provide your response in a clear, pedagogical manner, breaking down complex concepts for easier understanding. If appropriate, suggest areas where further research might be needed.
-After formulating your initial response, take a moment to reflect on your answer. Consider:
-1. Have you addressed all aspects of the question?
-2. Are there any potential biases or assumptions in your explanation?
-3. Is your explanation clear and accessible to someone with a general science background?
-4. Have you adequately conveyed the uncertainties or debates surrounding this topic?
-Based on this reflection, refine your answer as needed.
-Remember, while you have extensive knowledge, it's okay to acknowledge the limits of current scientific understanding. If parts of the question cannot be answered definitively, explain why.
-Relevant documents:{context}
-
-Question: {question}
-
-Initial Response:
-[Your initial response here]
-
-Reflection and Refinement:
-[Your reflections and any refinements to your answer here]
-
-Final Response:
-[Your final, refined answer here]"""
+    try:
+        loaders = []
 
-def make_rag_qa_answer(query, top_k = 10):
+        documents = []
+        my_bar = st.progress(0, text='adding documents to LLM context')
 
+        for i, row in papers_df.iterrows():
+            content = f"Paper {i+1}: {row['title']}\n{row['abstract']}\n\n"
+            metadata = {"source": row['ads_id']}
+            doc = Document(page_content=content, metadata=metadata)
+            documents.append(doc)
+            my_bar.progress((i+1)/len(papers_df), text='adding documents to LLM context')
 
-    # try:
-        absts, fhdrs, rs = Library2(query, top_k = top_k)
-    
-        temp_abst = ''
-        loaders = []
-        for i in range(len(absts)):
-            temp_abst = absts[i]
-    
-            try:
-                text_file = open("absts/"+fhdrs[i]+".txt", "w")
-            except:
-                os.mkdir('absts')
-                text_file = open("absts/"+fhdrs[i]+".txt", "w")
-            n = text_file.write(temp_abst)
-            text_file.close()
-            loader = TextLoader("absts/"+fhdrs[i]+".txt")
-            loaders.append(loader)
-    
         text_splitter = RecursiveCharacterTextSplitter(chunk_size=150, chunk_overlap=50, add_start_index=True)
-    
-        splits = text_splitter.split_documents([loader.load()[0] for loader in loaders])
+
+        splits = text_splitter.split_documents(documents)
         vectorstore = Chroma.from_documents(documents=splits, embedding=st.session_state.embeddings, collection_name='retdoc4')
         # retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 6,  "fetch_k": len(splits)})
         retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 6})
-    
-        for i in range(len(absts)):
-            try:
-                os.remove("absts/"+fhdrs[i]+".txt")
-            except:
-                print("absts/"+fhdrs[i]+".txt not found")
 
         if st.session_state.question_type == 'Bibliometric':
             template = bibliometric_prompt
@@ -713,79 +466,33 @@ def make_rag_qa_answer(query, top_k = 10):
         else:
             template = regular_prompt
         prompt = PromptTemplate.from_template(template)
-    
+
         def format_docs(docs):
             return "\n\n".join(doc.page_content for doc in docs)
-    
-    
+
         rag_chain_from_docs = (
             RunnablePassthrough.assign(context=(lambda x: format_docs(x["context"])))
             | prompt
             | st.session_state.gen_llm
             | StrOutputParser()
         )
-    
+
         rag_chain_with_source = RunnableParallel(
             {"context": retriever, "question": RunnablePassthrough()}
         ).assign(answer=rag_chain_from_docs)
-    
         rag_answer = rag_chain_with_source.invoke(query, )
-    
         vectorstore.delete_collection()
 
-    # except:
-    #     st.write('heavy load! please wait 10 seconds and try again.')
-    
-        return rag_answer, rs
+    except:
+        st.subheader('heavy load! please wait 10 seconds and try again.')
+
+    return rag_answer
 
 def guess_question_type(query: str):
-    categorization_prompt = """You are an expert astrophysicist and computer scientist specializing in linguistics and semantics. Your task is to categorize a given query into one of the following categories:
-
-    1. Summarization
-    2. Single-paper factual
-    3. Multi-paper factual
-    4. Named entity recognition
-    5. Jargon-specific questions / overloaded words
-    6. Time-sensitive
-    7. Consensus evaluation
-    8. What-ifs and counterfactuals
-    9. Compositional
-
-    Analyze the query carefully, considering its content, structure, and implications. Then, determine which of the above categories best fits the query.
-
-    In your analysis, consider the following:
-    - Does the query ask for a well-known datapoint or mechanism?
-    - Can it be answered by a single paper or does it require multiple sources?
-    - Does it involve proper nouns or specific scientific terms?
-    - Is it time-dependent or likely to change in the near future?
-    - Does it require evaluating consensus across multiple sources?
-    - Is it a hypothetical or counterfactual question?
-    - Does it need to be broken down into sub-queries (i.e. compositional)?
-
-    After your analysis, categorize the query into one of the nine categories listed above.
-
-    Provide a brief explanation for your categorization, highlighting the key aspects of the query that led to your decision.
-
-    Present your final answer in the following format:
-
-    <categorization>
-    Category: [Selected category]
-    Explanation: [Your explanation for the categorization]
-    </categorization>"""
-    # st.write('invoking hyde generation')
-
-    # message = self.generation_client.messages.create(
-    #         model = self.generation_model,
-    #         max_tokens = self.max_doclen,
-    #         temperature = self.temperature,
-    #         system = prompt,
-    #         messages=[{ "role": "user",
-    #                 "content": [{"type": "text", "text": query,}] }]
-    #     )
-    # return message.content[0].text
-
-    messages = [("system",categorization_prompt,),("human", query),]
-    return st.session_state.ec.generation_client.invoke(messages).content
+
+    gen_client = openai_llm(temperature=0,model_name='gpt-4o-mini', openai_api_key = openai_key)
+    messages = [("system",question_categorization_prompt,),("human", query),]
+    return gen_client.invoke(messages).content
 
 class OverallConsensusEvaluation(BaseModel):
     consensus: Literal["Strong Agreement", "Moderate Agreement", "Weak Agreement", "No Clear Consensus", "Weak Disagreement", "Moderate Disagreement", "Strong Disagreement"] = Field(
@@ -840,141 +547,83 @@ def evaluate_overall_consensus(query: str, abstracts: List[str]) -> OverallConse
 
     return response
 
-def create_embedding_plot(rs):
-    """
-    function to create embedding plot
-    """
-
-    pltsource = ColumnDataSource(data=dict(
-        x=st.session_state.umap_x,
-        y=st.session_state.umap_y,
-        title=st.session_state.titles,
-        link=st.session_state.bibcode,
-    ))
-
-    rsflag = np.zeros((len(st.session_state.ids),))
-    rsflag[np.array([k for k in rs])] = 1
-
-    # outflag = np.zeros((len(st.session_state.ids),))
-    # outflag[np.array([k for k in find_outliers(rs)])] = 1
-    pltsource.data['colors'] = rsflag * 0.8 + 0.1
-    # pltsource.data['colors'][outflag] = 0.5
-    pltsource.data['sizes'] = (rsflag + 1)**5 / 100
-
-    TOOLTIPS = """
-    <div style="width:300px;">
-    ID: $index
-    ($x, $y)
-    @title <br>
-    @link <br> <br>
-    </div>
-    """
-
-    mapper = linear_cmap(field_name="colors", palette=Spectral5, low=0., high=1.)
-
-    p = figure(width=700, height=900, tooltips=TOOLTIPS, x_range=(0, 20), y_range=(-4.2,18),
-            title="UMAP projection of embeddings for the astro-ph corpus")
 
-    p.axis.visible=False
-    p.grid.visible=False
-    p.outline_line_alpha = 0.
 
-    p.circle('x', 'y', radius='sizes', source=pltsource, alpha=0.3, fill_color=mapper, fill_alpha='colors', line_color="lightgrey",line_alpha=0.1)
 
-    return p
+# ---------------------------------------
 
-if submit_button:
-
-    keywords = [kw.strip() for kw in extra_keywords.split(',')] if extra_keywords else []
-    toggles = {'Keyword weighting': toggle_a, 'Time weighting': toggle_b, 'Citation weighting': toggle_c}
-    
-    if (method == "Semantic search"):
-        with st.spinner('set retrieval method to'+ method):
-            st.session_state.ec = EmbeddingRetrievalSystem()
-    elif (method == "Semantic search + HyDE"):
-        with st.spinner('set retrieval method to'+ method):
-            st.session_state.ec = HydeRetrievalSystem()
-    elif (method == "Semantic search + HyDE + CoHERE"):
-        with st.spinner('set retrieval method to'+ method):
-            st.session_state.ec = HydeCohereRetrievalSystem()
-    st.toast('loaded retrieval system')
-    
+if st.session_state.get('runpfdr'):
     with st.spinner(search_text_list[np.random.choice(len(search_text_list))]):
+        st.write('Settings: [Kw:',toggle_a, 'Time:',toggle_b, 'Cite:',toggle_c, '] top_k:',top_k, 'retrieval:',method)
 
-        st.session_state.ec.query_input_keywords = keywords
-        st.session_state.ec.toggles = toggles
-        st.session_state.ec.question_type = question_type
-        st.session_state.ec.rag_method = method
-        st.session_state.ec.gen_method = method2
-        
-        if method2 == "Basic RAG":
-            st.session_state.gen_method = 'rag'
-        elif method2 == "ReAct Agent":
-            st.session_state.gen_method = 'agent'
+        papers_df = run_query_ret(st.session_state.query)
+        st.header(st.session_state.query)
+        st.subheader('top-k relevant papers:')
+        st.data_editor(papers_df, column_config = {'ADS Link':st.column_config.LinkColumn(display_text= 'https://ui.adsabs.harvard.edu/abs/(.*?)/abstract')})
 
-        if st.session_state.gen_method == 'agent':
-            answer = run_agent_qa(query, top_k)
-            rs = get_topk(query, top_k)
+    with st.spinner(gen_text_list[np.random.choice(len(gen_text_list))]):
 
+        if st.session_state.gen_method == 'agent':
+            answer = run_agent_qa(st.session_state.query)
             answer_text = answer['output']
+            st.subheader('Answer with '+method2)
             st.write(answer_text)
-
             file_path = "agent_trace.txt"
             with open(file_path, 'r') as file:
                 intermediate_steps = file.read()
             st.expander('Intermediate steps', expanded=False).write(intermediate_steps)
 
         elif st.session_state.gen_method == 'rag':
-            answer, rs = make_rag_qa_answer(query, top_k)
+            answer = run_rag_qa(query, papers_df)
+            st.subheader('Answer with '+method2)
             answer_text = answer['answer']
             st.write(answer_text)
 
-        triggered_keywords = st.session_state.ec.query_kws
-
-        with st.spinner('compiling top-k papers'+ method):
-            papers_df = get_paper_df(rs)
-            
-            with st.expander("Relevant papers", expanded=True):
-                # st.dataframe(papers_df, hide_index=True)
-                st.data_editor(papers_df, column_config = {'ADS Link':st.column_config.LinkColumn(display_text= 'https://ui.adsabs.harvard.edu/abs/(.*?)/abstract')})
-
+        query_kws = get_keywords(query)
+        input_kws = st.session_state.keywords
+        query_kws = query_kws + input_kws
+        triggered_keywords = query_kws + input_kws
         st.write('**Triggered keywords:** `'+ "`, `".join(triggered_keywords)+'`')
-    
+
     col1, col2 = st.columns(2)
 
     with col1:
-        with st.expander("Evaluating question type", expanded=True):
-            st.subheader("Question type suggestion")
-            question_type_gen = guess_question_type(query)
-            if '<categorization>' in question_type_gen:
-                question_type_gen = question_type_gen.split('<categorization>')[1]
-            if '</categorization>' in question_type_gen:
-                question_type_gen = question_type_gen.split('</categorization>')[0]
-            question_type_gen = question_type_gen.replace('\n','  \n')
-            st.markdown(question_type_gen)
+        with st.spinner("Evaluating question type"):
+            with st.expander("Question type", expanded=True):
+                st.subheader("Question type suggestion")
+                question_type_gen = guess_question_type(query)
+                if '<categorization>' in question_type_gen:
+                    question_type_gen = question_type_gen.split('<categorization>')[1]
+                if '</categorization>' in question_type_gen:
+                    question_type_gen = question_type_gen.split('</categorization>')[0]
+                question_type_gen = question_type_gen.replace('\n','  \n')
+                st.markdown(question_type_gen)
 
     with col2:
-        with st.expander("Evaluating abstract consensus", expanded=True):
-            consensus_answer = evaluate_overall_consensus(query, [st.session_state.abstracts[i] for i in rs])
-            st.subheader("Consensus: "+consensus_answer.consensus)
-            st.markdown(consensus_answer.explanation)
-            st.markdown('Relevance of retrieved papers to answer: %.1f' %consensus_answer.relevance_score)
-                
+        with st.spinner("Evaluating abstract consensus"):
+            with st.expander("Abstract consensus", expanded=True):
+                consensus_answer = evaluate_overall_consensus(query, [papers_df['abstract'][i] for i in range(len(papers_df))])
+                st.subheader("Consensus: "+consensus_answer.consensus)
+                st.markdown(consensus_answer.explanation)
+                st.markdown('Relevance of retrieved papers to answer: %.1f' %consensus_answer.relevance_score)
+
     session_vars = {
         "runtime": "pathfinder_v1_online",
         "query": query,
         "question_type": question_type,
-        'Keyword weighting': toggle_a, 
-        'Time weighting': toggle_b, 
+        'Keyword weighting': toggle_a,
+        'Time weighting': toggle_b,
         'Citation weighting': toggle_c,
         "rag_method" : method,
         "gen_method" : method2,
         "answer" : answer_text,
-        "topk" : ['%.0f' %i for i in rs],
-        "topk_scores" : ['%.6f' %rs[i] for i in rs], 
+        "question_type": question_type_gen,
+        "consensus": consensus_answer.explanation,
+        "topk" : list(papers_df['ads_id']),
+        "topk_scores" : list(papers_df['Relevance']),
         "topk_papers": list(papers_df['ADS Link']),
     }
-    
+
     @st.fragment()
     def download_op(data):
         json_string = json.dumps(data)
@@ -983,12 +632,9 @@ if submit_button:
             file_name="pathfinder_data.json",
             mime="application/json",
             data=json_string,)
-        
-    with st.sidebar:
-        download_op(session_vars)
-        
-    # embedding_plot = create_embedding_plot(rs)
-    # st.bokeh_chart(embedding_plot)
-    
+
+    # with st.sidebar:
+    download_op(session_vars)
+
 else:
     st.info("Use the sidebar to tweak the search parameters to get better results.")