beacon / quick_pathfinder_query.py
kiyer's picture
Adopt Pathfinder keyword-weighted retrieval per quick_pathfinder_query.py
3e7f715
Raw
History Blame Contribute Delete
936 Bytes
from pathfinder_setup_fns import *
def get_df_query(query, top_k = 10, date_range = None, ec=ec):
input_keywords = ['']
query_keywords = get_keywords(query)
ec.query_input_keywords = input_keywords+query_keywords
ec.toggles = ['Keywords']
ec.hyde = False
ec.rerank = False
ec.date_range = date_range
rs, small_df = ec.retrieve(query, top_k = top_k, return_scores=True)
formatted_df = ec.return_formatted_df(rs, small_df)
qembd = np.array(ec.query_embedding)
return formatted_df, qembd
sfh_query = 'What can the star formation histories of distant galaxies obtained from SED fitting tell us about their evolution?'
start_date = datetime.date(1996, 1, 1)
# end_date = datetime.date(2012, 12, 31)
end_date = datetime.datetime.now().date()
date_range = (start_date, end_date)
sfh_df, qembd = get_df_query(sfh_query, date_range=date_range, ec=ec)
print(qembd.shape)
sfh_df[['title','date']]