hv68/sample_tool_1

1import streamlit as st2import pandas as pd3import sys4import os5from datasets import load_from_disk6# from st_aggrid import AgGrid, GridOptionsBuilder, GridUpdateMode7from sklearn.metrics.pairwise import cosine_similarity8import numpy as np9import time10from annotated_text import annotated_text11 12 13ABSOLUTE_PATH = os.path.dirname(__file__)14ASSETS_PATH = os.path.join(ABSOLUTE_PATH, 'model_assets')15 16 17from nltk.data import find18import nltk19import gensim20 21@st.cache(suppress_st_warning=True, allow_output_mutation=True)22def get_embed_model():23  nltk.download("word2vec_sample")24  word2vec_sample = str(find('models/word2vec_sample/pruned.word2vec.txt'))25 26  model = gensim.models.KeyedVectors.load_word2vec_format(word2vec_sample, binary=False)27  return model28 29@st.cache(suppress_st_warning=True, allow_output_mutation=True)30def get_top_n_closest(query_word, candidate, n):31  model = get_embed_model()32  t = time.time()33  p_c = preprocess_text(candidate)34  similarity = []35  t = time.time()36  for i in p_c:37    try:38      similarity.append(model.similarity(query_word, i))39    except:40      similarity.append(0)41  top_n = min(len(p_c), n)42  t = time.time()43  sorted = (-1*np.array(similarity)).argsort()[:top_n]44  top = [p_c[i] for i in sorted]45  return  top46 47@st.cache(suppress_st_warning=True, allow_output_mutation=True)48def annotate_text(text, words):49  annotated = [text]50  for word in words:51    for i in range(len(annotated)):52      if type(annotated[i]) != str:53        continue54      string =  annotated[i]55      try:56        index = string.index(word)57      except:58        continue59      first = string[:index]60      second = (string[index:index+len(word)],'SIMILAR')61      third = string[index+len(word):]62      annotated = annotated[:i]  + [first, second, third] + annotated[i+1:] 63  return tuple(annotated)64 65 66@st.cache(suppress_st_warning=True, allow_output_mutation=True)67def preprocess_text(s):68 return list(filter(lambda x: x!= '', (''.join(c if c.isalnum() or c == ' ' else ' ' for c in s)).split(' ')))69 70@st.cache(suppress_st_warning=True, allow_output_mutation=True)71def get_pairwise_distances(model):72	df = pd.read_csv(f"{ASSETS_PATH}/{model}/pairwise_distances.csv").set_index('index')73	return df74 75@st.cache(suppress_st_warning=True, allow_output_mutation=True)76def get_pairwise_distances_chunked(model, chunk):77	# for df in pd.read_csv(f"{ASSETS_PATH}/{model}/pairwise_distances.csv", chunksize = 16):78	# 	print(df.iloc[0]['queries'])79	# 	if chunk == int(df.iloc[0]['queries']):80	# 		return df81	return get_pairwise_distances(model)82@st.cache(suppress_st_warning=True, allow_output_mutation=True)83def get_query_strings():84	df = pd.read_json(f"{ASSETS_PATH}/IUR_Reddit_test_queries_english.jsonl", lines = True)85	df['index'] = df.reset_index().index86	return df87	# df['partition'] = df['index']%10088	# df.to_parquet(f"{ASSETS_PATH}/IUR_Reddit_test_queries_english.parquet", index = 'index', partition_cols = 'partition')89	90	# return pd.read_parquet(f"{ASSETS_PATH}/IUR_Reddit_test_queries_english.parquet", columns=['fullText', 'index', 'authorIDs'])91@st.cache(suppress_st_warning=True, allow_output_mutation=True)92def get_candidate_strings():93	df = pd.read_json(f"{ASSETS_PATH}/IUR_Reddit_test_candidates_english.jsonl", lines = True)94	df['i'] = df['index']95	df = df.set_index('i')96	# df['index'] = df.reset_index().index97 98	return df99	# df['partition'] = df['index']%100100	# df.to_parquet(f"{ASSETS_PATH}/IUR_Reddit_test_candidates_english.parquet", index = 'index', partition_cols = 'partition')101	# return pd.read_parquet(f"{ASSETS_PATH}/IUR_Reddit_test_candidates_english.parquet", columns=['fullText', 'index', 'authorIDs'])102@st.cache(suppress_st_warning=True, allow_output_mutation=True)103def get_embedding_dataset(model):104	data = load_from_disk(f"{ASSETS_PATH}/{model}/embedding")105	return data106@st.cache(suppress_st_warning=True, allow_output_mutation=True)107def get_bad_queries(model):108	df = get_query_strings().iloc[list(get_pairwise_distances(model)['queries'].unique())][['fullText', 'index', 'authorIDs']]109	return df110@st.cache(suppress_st_warning=True, allow_output_mutation=True)111def get_gt_candidates(model, author):112	gt_candidates = get_candidate_strings()113	df = gt_candidates[gt_candidates['authorIDs'] == author]114	return df115@st.cache(suppress_st_warning=True, allow_output_mutation=True)116def get_candidate_text(l):117	return get_candidate_strings().at[l,'fullText']118 119@st.cache(suppress_st_warning=True, allow_output_mutation=True)120def get_annotated_text(text, word, pos):121	print("here", word, pos)122	start= text.index(word, pos)123	end = start+len(word)124	return (text[:start], (text[start:end ], 'SELECTED'), text[end:]), end125 126# class AgGridBuilder:127# 	__static_key = 0128# 	def build_ag_grid(table, display_columns):129# 		AgGridBuilder.__static_key += 1130# 		options_builder = GridOptionsBuilder.from_dataframe(table[display_columns])131# 		options_builder.configure_pagination(paginationAutoPageSize=False, paginationPageSize=10)132# 		options_builder.configure_selection(selection_mode= 'single', pre_selected_rows = [0])133# 		options = options_builder.build()134# 		return AgGrid(table, gridOptions = options, fit_columns_on_grid_load=True, key = AgGridBuilder.__static_key, reload_data = True, update_mode = GridUpdateMode.SELECTION_CHANGED | GridUpdateMode.VALUE_CHANGED)135 136if __name__ == '__main__':137	st.set_page_config(layout="wide")138 139	models = filter(lambda file_name: os.path.isdir(f"{ASSETS_PATH}/{file_name}") and not file_name.endswith(".parquet"), os.listdir(ASSETS_PATH))140 141	with st.sidebar:142		current_model = st.selectbox(143			"Select Model to analyze", 144			models145			)146 147	pairwise_distances = get_pairwise_distances(current_model)148	embedding_dataset = get_embedding_dataset(current_model)149 150	candidate_string_grid = None151	gt_candidate_string_grid  = None152	with st.container():153		t1 = time.time()154		st.title("Full Text")155		col1, col2 = st.columns([14, 2])156		t2 = time.time()157		query_table = get_bad_queries(current_model)158		t3 = time.time()159		print(query_table)160		with col2:161			index = st.number_input('Enter Query number to inspect', min_value = 0, max_value = query_table.shape[0], step = 1)162			query_text = query_table.loc[index]['fullText']163			preprocessed_query_text = preprocess_text(query_text)164			text_highlight_index = st.number_input('Enter word #', min_value = 0, max_value = len(preprocessed_query_text), step = 1)165			query_index = int(query_table.iloc[index]['index'])166 167		with col1:168			if 'pos_highlight' not in st.session_state or text_highlight_index == 0:169				st.session_state['pos_highlight'] = text_highlight_index170				st.session_state['pos_history'] = [0]171 172			if st.session_state['pos_highlight'] > text_highlight_index:173				st.session_state['pos_history'] = st.session_state['pos_history'][:-2]174				if len(st.session_state['pos_history']) == 0:175					st.session_state['pos_history'] = [0]176			print("pos", st.session_state['pos_history'], st.session_state['pos_highlight'], text_highlight_index)177			anotated_text_, pos = get_annotated_text(query_text, preprocessed_query_text[text_highlight_index-1], st.session_state['pos_history'][-1]) if text_highlight_index >= 1 else ((query_text), 0)178			if st.session_state['pos_highlight'] < text_highlight_index:179				st.session_state['pos_history'].append(pos)180			st.session_state['pos_highlight'] = text_highlight_index181			annotated_text(*anotated_text_)182			# annotated_text("Lol, this" , ('guy', 'SELECTED') , "is such a PR chameleon. \n\n In the Chan Zuckerberg Initiative announcement, he made it sound like he was giving away all his money to charity <PERSON> or <PERSON>. http://www.businessinsider.in/Mark-Zuckerberg-says-hes-giving-99-of-his-Facebook-shares-45-billion-to-charity/articleshow/50005321.cms Apparently, its just a VC fund. And there are still people out there who believe Facebook.org was an initiative to bring Internet to the poor.")183			t4 = time.time()184 185		print(f"query time query text: {t3-t2}, total time: {t4-t1}")186	with st.container():187		st.title("Top 16 Recommended Candidates")188		col1, col2, col3 = st.columns([10, 4, 2])	189		rec_candidates = pairwise_distances[pairwise_distances["queries"]==query_index]['candidates']190		print(rec_candidates)191		l = list(rec_candidates)192		with col3:193			candidate_rec_index = st.number_input('Enter recommended candidate number to inspect', min_value = 0, max_value = len(l), step = 1)194			print("l:",l, query_index)195			pairwise_candidate_index = int(l[candidate_rec_index])196		with col1:197			st.header("Text")198			t1 = time.time()199			candidate_text = get_candidate_text(pairwise_candidate_index)200 201			if st.session_state['pos_highlight'] == 0:202				annotated_text(candidate_text)203			else:204				top_n_words_to_highlight = get_top_n_closest(preprocessed_query_text[text_highlight_index-1], candidate_text, 4)205				print("TOPN", top_n_words_to_highlight)206				annotated_text(*annotate_text(candidate_text, top_n_words_to_highlight))207 208			t2 = time.time()209		with col2:210			st.header("Cosine Distance")211			st.write(float(pairwise_distances[\212				( pairwise_distances['queries'] == query_index ) \213				&214				( pairwise_distances['candidates'] == pairwise_candidate_index)]['distances']))215		print(f"candidate string retreival: {t2-t1}")216	with st.container():217		t1 = time.time()218		st.title("Candidates With Same Authors As Query")219		col1, col2, col3 = st.columns([10, 4, 2])220		t2 = time.time()221		gt_candidates = get_gt_candidates(current_model, query_table.iloc[query_index]['authorIDs'][0])222		t3 = time.time()223 224		with col3:225			candidate_index = st.number_input('Enter ground truthnumber to inspect', min_value = 0, max_value = gt_candidates.shape[0], step = 1)226			print(gt_candidates.head())227			gt_candidate_index = int(gt_candidates.iloc[candidate_index]['index'])228		with col1:229			st.header("Text")230			st.write(gt_candidates.iloc[candidate_index]['fullText'])231		with col2:232			t4 = time.time()233			st.header("Cosine Distance")234			indices = list(embedding_dataset['candidates']['index'])235			st.write(1-cosine_similarity(np.array([embedding_dataset['queries'][query_index]['embedding']]), np.array([embedding_dataset['candidates'][indices.index(gt_candidate_index)]['embedding']]))[0,0])236			t5 = time.time()237		print(f"find gt candidates: {t3-t2}, find cosine: {t5-t4}, total: {t5-t1}")238