fix

app.py

@@ -13,23 +13,73 @@ from annotated_text import annotated_text
 ABSOLUTE_PATH = os.path.dirname(__file__)
 ASSETS_PATH = os.path.join(ABSOLUTE_PATH, 'model_assets')
 
-@st.cache(suppress_st_warning=True, allow_output_mutation=True)
+
+from nltk.data import find
+import nltk
+import gensim
+
+@st.cache_data
+def get_embed_model():
+  nltk.download("word2vec_sample")
+  word2vec_sample = str(find('models/word2vec_sample/pruned.word2vec.txt'))
+ 
+  model = gensim.models.KeyedVectors.load_word2vec_format(word2vec_sample, binary=False)
+  return model
+
+@st.cache_data
+def get_top_n_closest(query_word, candidate, n):
+  model = get_embed_model()
+  t = time.time()
+  p_c = preprocess_text(candidate)
+  similarity = []
+  t = time.time()
+  for i in p_c:
+    try:
+      similarity.append(model.similarity(query_word, i))
+    except:
+      similarity.append(0)
+  top_n = min(len(p_c), n)
+  t = time.time()
+  sorted = (-1*np.array(similarity)).argsort()[:top_n]
+  top = [p_c[i] for i in sorted]
+  return  top
+
+@st.cache_data
+def annotate_text(text, words):
+  annotated = [text]
+  for word in words:
+    for i in range(len(annotated)):
+      if type(annotated[i]) != str:
+        continue
+      string =  annotated[i]
+      try:
+        index = string.index(word)
+      except:
+        continue
+      first = string[:index]
+      second = (string[index:index+len(word)],'SIMILAR')
+      third = string[index+len(word):]
+      annotated = annotated[:i]  + [first, second, third] + annotated[i+1:] 
+  return tuple(annotated)
+
+
+@st.cache_data
 def preprocess_text(s):
  return list(filter(lambda x: x!= '', (''.join(c if c.isalnum() or c == ' ' else ' ' for c in s)).split(' ')))
 
-@st.cache(suppress_st_warning=True, allow_output_mutation=True)
+@st.cache_data
 def get_pairwise_distances(model):
 	df = pd.read_csv(f"{ASSETS_PATH}/{model}/pairwise_distances.csv").set_index('index')
 	return df
 
-@st.cache(suppress_st_warning=True, allow_output_mutation=True)
+@st.cache_data
 def get_pairwise_distances_chunked(model, chunk):
 	# for df in pd.read_csv(f"{ASSETS_PATH}/{model}/pairwise_distances.csv", chunksize = 16):
 	# 	print(df.iloc[0]['queries'])
 	# 	if chunk == int(df.iloc[0]['queries']):
 	# 		return df
 	return get_pairwise_distances(model)
-@st.cache(suppress_st_warning=True, allow_output_mutation=True)
+@st.cache_data
 def get_query_strings():
 	df = pd.read_json(f"{ASSETS_PATH}/IUR_Reddit_test_queries_english.jsonl", lines = True)
 	df['index'] = df.reset_index().index
@@ -38,7 +88,7 @@ def get_query_strings():
 	# df.to_parquet(f"{ASSETS_PATH}/IUR_Reddit_test_queries_english.parquet", index = 'index', partition_cols = 'partition')
 	
 	# return pd.read_parquet(f"{ASSETS_PATH}/IUR_Reddit_test_queries_english.parquet", columns=['fullText', 'index', 'authorIDs'])
-@st.cache(suppress_st_warning=True, allow_output_mutation=True)
+@st.cache_data
 def get_candidate_strings():
 	df = pd.read_json(f"{ASSETS_PATH}/IUR_Reddit_test_candidates_english.jsonl", lines = True)
 	df['i'] = df['index']
@@ -49,24 +99,24 @@ def get_candidate_strings():
 	# df['partition'] = df['index']%100
 	# df.to_parquet(f"{ASSETS_PATH}/IUR_Reddit_test_candidates_english.parquet", index = 'index', partition_cols = 'partition')
 	# return pd.read_parquet(f"{ASSETS_PATH}/IUR_Reddit_test_candidates_english.parquet", columns=['fullText', 'index', 'authorIDs'])
-@st.cache(suppress_st_warning=True, allow_output_mutation=True)
+@st.cache_data
 def get_embedding_dataset(model):
 	data = load_from_disk(f"{ASSETS_PATH}/{model}/embedding")
 	return data
-@st.cache(suppress_st_warning=True, allow_output_mutation=True)
+@st.cache_data
 def get_bad_queries(model):
 	df = get_query_strings().iloc[list(get_pairwise_distances(model)['queries'].unique())][['fullText', 'index', 'authorIDs']]
 	return df
-@st.cache(suppress_st_warning=True, allow_output_mutation=True)
+@st.cache_data
 def get_gt_candidates(model, author):
 	gt_candidates = get_candidate_strings()
 	df = gt_candidates[gt_candidates['authorIDs'] == author]
 	return df
-@st.cache(suppress_st_warning=True, allow_output_mutation=True)
+@st.cache_data
 def get_candidate_text(l):
 	return get_candidate_strings().at[l,'fullText']
 
-@st.cache(suppress_st_warning=True, allow_output_mutation=True)
+@st.cache_data
 def get_annotated_text(text, word, pos):
 	print("here", word, pos)
 	start= text.index(word, pos)
@@ -146,7 +196,15 @@ if __name__ == '__main__':
 		with col1:
 			st.header("Text")
 			t1 = time.time()
-			st.write(get_candidate_text(pairwise_candidate_index))
+			candidate_text = get_candidate_text(pairwise_candidate_index)
+
+			if st.session_state['pos_highlight'] == 0:
+				annotated_text(candidate_text)
+			else:
+				top_n_words_to_highlight = get_top_n_closest(preprocessed_query_text[text_highlight_index-1], candidate_text, 4)
+				print("TOPN", top_n_words_to_highlight)
+				annotated_text(*annotate_text(candidate_text, top_n_words_to_highlight))
+
 			t2 = time.time()
 		with col2:
 			st.header("Cosine Distance")

requirements.txt

@@ -1,4 +1,7 @@
 scikit-learn==1.2.0
 numpy==1.23.5
 pandas==1.5.2
-st-annotated-text==3.0.0
+st-annotated-text==3.0.0
+nltk==3.8.1
+gensim==4.3.1
+streamlit==1.20.0