Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| st.set_page_config(layout="wide") | |
| import numpy as np | |
| from abc import ABC, abstractmethod | |
| from typing import List, Dict, Any, Tuple | |
| from collections import defaultdict | |
| from tqdm import tqdm | |
| import pandas as pd | |
| from datetime import datetime, date | |
| from datasets import load_dataset, load_from_disk | |
| from collections import Counter | |
| import yaml, json, requests, sys, os, time | |
| import concurrent.futures | |
| from langchain import hub | |
| from langchain_openai import ChatOpenAI as openai_llm | |
| from langchain_openai import OpenAIEmbeddings | |
| from langchain_core.runnables import RunnableConfig, RunnablePassthrough, RunnableParallel | |
| from langchain_core.prompts import PromptTemplate | |
| from langchain_community.callbacks import StreamlitCallbackHandler | |
| from langchain_community.utilities import DuckDuckGoSearchAPIWrapper | |
| from langchain_community.vectorstores import Chroma | |
| from langchain_community.document_loaders import TextLoader | |
| from langchain.agents import create_react_agent, Tool, AgentExecutor | |
| from langchain.text_splitter import RecursiveCharacterTextSplitter | |
| from langchain_core.output_parsers import StrOutputParser | |
| from langchain.callbacks import FileCallbackHandler | |
| from langchain.callbacks.manager import CallbackManager | |
| import instructor | |
| from pydantic import BaseModel, Field | |
| from typing import List, Literal | |
| from nltk.corpus import stopwords | |
| import nltk | |
| from openai import OpenAI | |
| # import anthropic | |
| import cohere | |
| import faiss | |
| import spacy | |
| from string import punctuation | |
| import pytextrank | |
| nlp = spacy.load("en_core_web_sm") | |
| nlp.add_pipe("textrank") | |
| try: | |
| stopwords.words('english') | |
| except: | |
| nltk.download('stopwords') | |
| stopwords.words('english') | |
| 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 | |
| ts = time.time() | |
| # anthropic_key = st.secrets["anthropic_key"] | |
| openai_key = st.secrets["openai_key"] | |
| cohere_key = st.secrets['cohere_key'] | |
| gen_llm = openai_llm(temperature=0,model_name='gpt-4o-mini', openai_api_key = openai_key) | |
| consensus_client = instructor.patch(OpenAI(api_key=openai_key)) | |
| embed_client = OpenAI(api_key = openai_key) | |
| embed_model = "text-embedding-3-small" | |
| embeddings = OpenAIEmbeddings(model = embed_model, api_key = openai_key) | |
| st.image('local_files/pathfinder_logo.png') | |
| st.expander("About", 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). | |
| """ | |
| ) | |
| # ---------------- get data and set up session state --------------------------- | |
| if 'arxiv_corpus' not in st.session_state: | |
| with st.spinner('loading data...'): | |
| try: | |
| 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') | |
| else: | |
| arxiv_corpus = st.session_state.arxiv_corpus | |
| if 'ids' not in st.session_state: | |
| st.session_state.ids = arxiv_corpus['ads_id'] | |
| st.session_state.titles = arxiv_corpus['title'] | |
| st.session_state.abstracts = arxiv_corpus['abstract'] | |
| 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)) | |
| #--------------------------------------------------------------- | |
| # A hack to "clear" the previous result when submitting a new prompt. This avoids | |
| # the "previous run's text is grayed-out but visible during rerun" Streamlit behavior. | |
| class DirtyState: | |
| NOT_DIRTY = "NOT_DIRTY" | |
| DIRTY = "DIRTY" | |
| UNHANDLED_SUBMIT = "UNHANDLED_SUBMIT" | |
| def get_dirty_state() -> str: | |
| return st.session_state.get("dirty_state", DirtyState.NOT_DIRTY) | |
| def set_dirty_state(state: str) -> None: | |
| st.session_state["dirty_state"] = state | |
| def with_clear_container(submit_clicked: bool) -> bool: | |
| if get_dirty_state() == DirtyState.DIRTY: | |
| if submit_clicked: | |
| set_dirty_state(DirtyState.UNHANDLED_SUBMIT) | |
| st.experimental_rerun() | |
| else: | |
| set_dirty_state(DirtyState.NOT_DIRTY) | |
| if submit_clicked or get_dirty_state() == DirtyState.UNHANDLED_SUBMIT: | |
| set_dirty_state(DirtyState.DIRTY) | |
| return True | |
| return False | |
| # ---------------- define embedding retrieval systems -------------------------- | |
| def get_keywords(text): | |
| result = [] | |
| pos_tag = ['PROPN', 'ADJ', 'NOUN'] | |
| doc = nlp(text.lower()) | |
| for token in doc: | |
| if(token.text in nlp.Defaults.stop_words or token.text in punctuation): | |
| continue | |
| if(token.pos_ in pos_tag): | |
| result.append(token.text) | |
| return result | |
| def parse_doc(text, nret = 10): | |
| local_kws = [] | |
| doc = 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 | |
| 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.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] | |
| def make_embedding(self, text): | |
| str_embed = self.client.embeddings.create(input = [text], model = self.embed_model).data[0].embedding | |
| return str_embed | |
| def embed_batch(self, texts: List[str]) -> List[np.ndarray]: | |
| embeddings = self.client.embeddings.create(input=texts, model=self.embed_model).data | |
| return [np.array(embedding.embedding, dtype=np.float32) for embedding in embeddings] | |
| 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') | |
| # 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) | |
| 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) | |
| 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] | |
| kw_weight = np.zeros((len(topk_indices),)) + 0.1 | |
| for k in query_kws: | |
| for i in (range(len(topk_indices))): | |
| for j in range(len(sub_kws[i])): | |
| if k.lower() in sub_kws[i][j].lower(): | |
| kw_weight[i] = kw_weight[i] + 0.1 | |
| # print(i, k, sub_kws[i][j]) | |
| # kw_weight = kw_weight**0.36 / np.amax(kw_weight**0.36) | |
| kw_weight = kw_weight / np.amax(kw_weight) | |
| else: | |
| kw_weight = np.ones((len(topk_indices),)) | |
| if self.weight_date == True: | |
| sub_dates = [self.years[i] for i in topk_indices] | |
| 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 | |
| age_weight = (1 + np.exp(date_diff/0.7))**(-1) | |
| age_weight = age_weight / np.amax(age_weight) | |
| else: | |
| age_weight = np.ones((len(topk_indices),)) | |
| if self.weight_citation == True: | |
| # st.write('weighting by citations') | |
| sub_cites = np.array([self.cites[i] for i in topk_indices]) | |
| temp = sub_cites.copy() | |
| temp[sub_cites > 300] = 300. | |
| cite_weight = (1 + np.exp((300-temp)/42.0))**(-1.) | |
| cite_weight = cite_weight / np.amax(cite_weight) | |
| else: | |
| cite_weight = np.ones((len(topk_indices),)) | |
| similarities = similarities * (kw_weight) * (age_weight) * (cite_weight) | |
| # if self.weight_keywords: | |
| # keyword_matches = self.keyword_filter.filter(query) | |
| # self.query_kws = keyword_matches | |
| # kw_indices = np.zeros_like(similarities) | |
| # for s in keyword_matches: | |
| # if self.id_to_index[s] in topk_indices: | |
| # # print('yes', self.id_to_index[s], topk_indices[np.where(topk_indices == self.id_to_index[s])[0]]) | |
| # similarities[np.where(topk_indices == self.id_to_index[s])[0]] = similarities[np.where(topk_indices == self.id_to_index[s])[0]] * 10. | |
| # similarities = similarities / 10. | |
| 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] | |
| if return_scores: | |
| return {doc[0]: doc[1] for doc in top_results} | |
| # 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) | |
| self.generation_client = openai_llm(temperature=0,model_name='gpt-4o-mini', openai_api_key = openai_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 | |
| 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 | |
| def generate_docs(self, query: str): | |
| docs = [] | |
| for i in range(self.generate_n): | |
| # st.write('invoking hyde generation2') | |
| docs.append(self.generate_doc(query)) | |
| # with concurrent.futures.ThreadPoolExecutor() as executor: | |
| # st.write('invoking hyde generation2') | |
| # future_to_query = {executor.submit(self.generate_doc, query): query for i in range(self.generate_n)} | |
| # for future in concurrent.futures.as_completed(future_to_query): | |
| # query = future_to_query[future] | |
| # try: | |
| # data = future.result() | |
| # docs.append(data) | |
| # except Exception as exc: | |
| # pass | |
| return docs | |
| def embed_docs(self, docs: List[str]): | |
| return self.embed_batch(docs) | |
| class HydeCohereRetrievalSystem(HydeRetrievalSystem): | |
| def __init__(self, **kwargs): | |
| super().__init__(**kwargs) | |
| self.cohere_key = cohere_key | |
| self.cohere_client = cohere.Client(self.cohere_key) | |
| 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]]: | |
| 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} | |
| top_results = super().retrieve(query, top_k = rerank_top_k, time_result = time_result) | |
| # 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] | |
| 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]) | |
| if reweight: | |
| if time_result['has_temporal_aspect']: | |
| final_results = self.date_filter.filter(final_results, time_score = time_result['expected_recency_weight']) | |
| if self.weight_citation: self.citation_filter.filter(final_results) | |
| if return_scores: | |
| return {result[0]: result[2] for result in final_results} | |
| return [doc[0] for doc in final_results] | |
| def embed_docs(self, docs: List[str]): | |
| return self.embed_batch(docs) | |
| # ---------------------------------------------------------------- | |
| if 'ec' not in st.session_state: | |
| ec = HydeCohereRetrievalSystem(weight_keywords=True) | |
| st.session_state.ec = ec | |
| st.toast('loaded retrieval system') | |
| else: | |
| ec = st.session_state.ec | |
| def get_topk(query, top_k): | |
| print('running retrieval') | |
| rs = st.session_state.ec.retrieve(query, top_k, return_scores=True) | |
| return rs | |
| 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.bibcode[i][0:4] + ') ' + st.session_state.titles[i] + '\n' + st.session_state.abstracts[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 = [], [], [], [], [], [] | |
| 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]) | |
| kws.append(st.session_state.ads_kws[i]) | |
| return pd.DataFrame({ | |
| 'Title': papers, | |
| 'Relevance': scores, | |
| 'Year': yrs, | |
| 'ADS Link': links, | |
| 'Citations': cites, | |
| 'Keywords': kws, | |
| }) | |
| # def find_outliers(inp_simids, arxiv_cutoff_distance = 0.8): | |
| # | |
| # inp_simids = np.array(inp_simids) | |
| # | |
| # # Calculate the centroid for each point, excluding itself | |
| # orange_black_points = st.session_state.embed[inp_simids] | |
| # | |
| # topk_dists = [] | |
| # for i, point in enumerate(orange_black_points): | |
| # # Exclude the current point | |
| # other_points = np.delete(orange_black_points, i, axis=0) | |
| # # Calculate centroid of other points | |
| # centroid = np.mean(other_points, axis=0) | |
| # # Calculate distance from the point to this centroid | |
| # dist = np.sqrt(np.sum((point - centroid)**2)) | |
| # topk_dists.append(dist) | |
| # | |
| # topk_dists = np.array(topk_dists) | |
| # | |
| # # Separate distances for orange and black points | |
| # orange_distances = topk_dists[:len(inp_simids)] | |
| # black_distances = topk_dists[len(inp_simids):] | |
| # | |
| # # Calculate the median of distances | |
| # orange_black_distances = topk_dists | |
| # median_topk_distance = np.median(orange_black_distances) | |
| # | |
| # # def get_sims_and_dists(inp_data): | |
| # | |
| # # all_sims, all_dists = [], [] | |
| # | |
| # # np.random.seed(12) | |
| # # rand_indices = np.random.choice(inp_data.shape[0], size=return_n, replace=False) | |
| # | |
| # # for j in tqdm(range(len(rand_indices))): | |
| # | |
| # # i = rand_indices[j] | |
| # # inferred_vector = inp_data[i,:] | |
| # # sims, dists = find_closest_dists(i, inp_data, return_n + 1) | |
| # # all_sims.append(sims[1:]) | |
| # # all_dists.append(dists[1:]) | |
| # | |
| # # return np.array(all_sims), np.array(all_dists) | |
| # | |
| # # # Identify papers with distances greater than the 95th percentile | |
| # # _, all_dists = get_sims_and_dists(arxiv_ada_embeddings) | |
| # # arxiv_cutoff_distance = find_cutoff_dist(all_dists) | |
| # # hardcoding for now | |
| # outlier_indices = inp_simids[np.where(orange_black_distances > arxiv_cutoff_distance)[0]] | |
| # # outlier_titles = [titles[i] for i in outlier_indices] | |
| # | |
| # return outlier_indices #, outlier_titles | |
| 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 | |
| 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 | |
| search = DuckDuckGoSearchAPIWrapper() | |
| tools = [ | |
| Tool( | |
| name="Library", | |
| func=Library, | |
| description="A source of information pertinent to your question. Do not answer a question without consulting this!" | |
| ), | |
| Tool( | |
| name="Search", | |
| func=search.run, | |
| description="useful for when you need to look up knowledge about common topics or current events", | |
| ) | |
| ] | |
| 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 | |
| file_path = "agent_trace.txt" | |
| try: | |
| os.remove(file_path) | |
| except: | |
| pass | |
| 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) | |
| agent = create_react_agent(llm=gen_llm, tools=tools, prompt=prompt) | |
| st.session_state.agent = agent | |
| if 'agent_executor' not in st.session_state: | |
| agent_executor = AgentExecutor(agent=st.session_state.agent, tools=st.session_state.tools, verbose=True, handle_parsing_errors=True, callbacks=CallbackManager([file_handler])) | |
| st.session_state.agent_executor = agent_executor | |
| answer = st.session_state.agent_executor.invoke({"input": query,}) | |
| return answer | |
| def make_rag_qa_answer(query, top_k = 10): | |
| 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]) | |
| vectorstore = Chroma.from_documents(documents=splits, embedding=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)): | |
| os.remove("absts/"+fhdrs[i]+".txt") | |
| 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. | |
| 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:""" | |
| 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 | |
| | 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() | |
| return rag_answer, rs | |
| 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 | |
| class OverallConsensusEvaluation(BaseModel): | |
| consensus: Literal["Strong Agreement", "Moderate Agreement", "Weak Agreement", "No Clear Consensus", "Weak Disagreement", "Moderate Disagreement", "Strong Disagreement"] = Field( | |
| ..., | |
| description="The overall level of consensus between the query and the abstracts" | |
| ) | |
| explanation: str = Field( | |
| ..., | |
| description="A detailed explanation of the consensus evaluation" | |
| ) | |
| relevance_score: float = Field( | |
| ..., | |
| description="A score from 0 to 1 indicating how relevant the abstracts are to the query overall", | |
| ge=0, | |
| le=1 | |
| ) | |
| def evaluate_overall_consensus(query: str, abstracts: List[str]) -> OverallConsensusEvaluation: | |
| """ | |
| Evaluates the overall consensus of the abstracts in relation to the query in a single LLM call. | |
| """ | |
| prompt = f""" | |
| Query: {query} | |
| You will be provided with {len(abstracts)} scientific abstracts. Your task is to: | |
| 1. Evaluate the overall consensus between the query and the abstracts. | |
| 2. Provide a detailed explanation of your consensus evaluation. | |
| 3. Assign an overall relevance score from 0 to 1, where 0 means completely irrelevant and 1 means highly relevant. | |
| For the consensus evaluation, use one of the following levels: | |
| Strong Agreement, Moderate Agreement, Weak Agreement, No Clear Consensus, Weak Disagreement, Moderate Disagreement, Strong Disagreement | |
| Here are the abstracts: | |
| {' '.join([f"Abstract {i+1}: {abstract}" for i, abstract in enumerate(abstracts)])} | |
| Provide your evaluation in a structured format. | |
| """ | |
| response = consensus_client.chat.completions.create( | |
| model="gpt-4o-mini", # used to be "gpt-4", | |
| response_model=OverallConsensusEvaluation, | |
| messages=[ | |
| {"role": "system", "content": """You are an assistant with expertise in astrophysics for question-answering tasks. | |
| Evaluate the overall consensus of the retrieved scientific abstracts in relation to a given query. | |
| If you don't know the answer, just say that you don't know. | |
| Use six sentences maximum and keep the answer concise."""}, | |
| {"role": "user", "content": prompt} | |
| ], | |
| temperature=0 | |
| ) | |
| return response | |
| # Streamlit app | |
| def main(): | |
| # st.title("Question Answering App") | |
| # Sidebar (Inputs) | |
| 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) | |
| if (method == "Semantic search"): | |
| with st.spinner('set retrieval method to'+ method): | |
| st.session_state.ec = EmbeddingRetrievalSystem(weight_keywords=True) | |
| elif (method == "Semantic search + HyDE"): | |
| with st.spinner('set retrieval method to'+ method): | |
| st.session_state.ec = HydeRetrievalSystem(weight_keywords=True) | |
| elif (method == "Semantic search + HyDE + CoHERE"): | |
| with st.spinner('set retrieval method to'+ method): | |
| st.session_state.ec = HydeCohereRetrievalSystem(weight_keywords=True) | |
| method2 = st.sidebar.radio("Generation complexity:", ["Basic RAG","ReAct Agent"]) | |
| 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("Select question type:", ["Single paper", "Multi-paper", "Summary"]) | |
| store_output = st.sidebar.button("Save output") | |
| # Main page (Outputs) | |
| # st.markdown(""" | |
| # <style> | |
| # .stTextInput > div > div { font-size: 50px; } | |
| # </style> | |
| # """, unsafe_allow_html=True) | |
| # st.markdown( | |
| # """ | |
| # <style> | |
| # textarea { | |
| # font-size: 3rem !important; | |
| # font-weight: bold; | |
| # font-family: "Times New Roman", Times, serif; | |
| # } | |
| # input { | |
| # font-size: 3rem !important; | |
| # font-weight: bold; | |
| # font-family: "Times New Roman", Times, serif; | |
| # } | |
| # </style> | |
| # """, | |
| # unsafe_allow_html=True, | |
| # ) | |
| # query = st.text_area("Ask me anything:", height=30) | |
| query = st.text_input("Ask me anything:") | |
| submit_button = st.button("Submit") | |
| if submit_button: | |
| 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...'] | |
| with st.spinner(search_text_list[np.random.choice(len(search_text_list))]): | |
| # Process inputs | |
| 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} | |
| # Generate outputs | |
| 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 | |
| # Display outputs | |
| if st.session_state.gen_method == 'agent': | |
| answer = run_agent_qa(query, top_k) | |
| rs = get_topk(query, top_k) | |
| st.write(answer["output"]) | |
| 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) | |
| st.write(answer['answer']) | |
| papers_df = get_paper_df(rs) | |
| embedding_plot = create_embedding_plot(rs) | |
| triggered_keywords = st.session_state.ec.query_kws | |
| st.write('**Triggered keywords:** `'+ "`, `".join(triggered_keywords)+'`') | |
| # consensus = estimate_consensus() | |
| 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')} | |
| ) | |
| # with st.expander("Embedding map", expanded=False): | |
| st.bokeh_chart(embedding_plot) | |
| col1, col2 = st.columns(2) | |
| with col1: | |
| 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: | |
| # st.subheader("Triggered Keywords") | |
| # st.write(", ".join(triggered_keywords)) | |
| 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) | |
| # st.write(f"{consensus:.2%}") | |
| else: | |
| st.info("Use the sidebar to tweak the search parameters to get better results.") | |
| if store_output: | |
| st.toast("Output stored successfully!") | |
| if __name__ == "__main__": | |
| main() | |