Update app.py

app.py

@@ -17,6 +17,11 @@ tokenizer_finbert = AutoTokenizer.from_pretrained("ProsusAI/finbert")
 kp_dict_finbert_checkpoint = "kp_dict_finbert.pickle"
 kp_cosine_finbert_checkpoint = "cosine_kp_finbert.pickle"
 
+tokenizer_sapbert = AutoTokenizer.from_pretrained("cambridgeltl/SapBERT-from-PubMedBERT-fulltext")
+model_sapbert = AutoModel.from_pretrained("cambridgeltl/SapBERT-from-PubMedBERT-fulltext", output_hidden_states=True)
+kp_dict_sapbert_checkpoint = "kp_dict_sapbert.pickle"
+kp_cosine_sapbert_checkpoint = "cosine_kp_sapbert.pickle"
+
 text = st.text_input("Enter word or key-phrase")
 exclude_words = st.radio("exclude_words",[True,False], help="Exclude results that contain any words in the query")
 exclude_text = st.radio("exclude_text",[True,False], help="Exclude results that contain the query (i.e exclude 'tomato soup recipe' if the query is 'tomato soup')")
@@ -28,7 +33,7 @@ with st.sidebar:
     k_diversify = st.number_input("Set of key-phrases to diversify from",10,30,20)
 
 #columns
-col1, col2 = st.columns(2)
+col1, col2, col3 = st.columns(3)
 #load kp dicts
 with open(kp_dict_checkpoint,'rb') as handle:
   kp_dict = pickle.load(handle)
@@ -38,11 +43,17 @@ with open(kp_dict_finbert_checkpoint,'rb') as handle:
   kp_dict_finbert = pickle.load(handle)
 keys_finbert = list(kp_dict_finbert.keys())
 
+with open(kp_dict_sapbert_checkpoint,'rb') as handle:
+  kp_dict_sapbert = pickle.load(handle)
+keys_sapbert = list(kp_dict_sapbert.keys())
+
 #load cosine distances of kp dict
 with open(kp_cosine_checkpoint,'rb') as handle:
   cosine_kp = pickle.load(handle)
 with open(kp_cosine_finbert_checkpoint,'rb') as handle:
   cosine_finbert_kp = pickle.load(handle)
+with open(kp_cosine_sapbert_checkpoint,'rb') as handle:
+  cosine_sapbert_kp = pickle.load(handle)
   
 def calculate_top_k(out, tokens,text,kp_dict,exclude_text=False,exclude_words=False, k=5):
   sim_dict = {}
@@ -100,11 +111,15 @@ if text:
   new_tokens.pop("KPS")
   new_tokens_finbert = concat_tokens([text], tokenizer_finbert)
   new_tokens_finbert.pop("KPS")
+  new_tokens_sapbert = concat_tokens([text], tokenizer_sapbert)
+  new_tokens_sapbert.pop("KPS")
   with torch.no_grad():
     outputs = model(**new_tokens)
     outputs_finbert = model_finbert(**new_tokens_finbert)
+    outputs_sapbert = model_sapbert(**new_tokens_sapbert)
   sim_dict = calculate_top_k(outputs, new_tokens, text, kp_dict, exclude_text=exclude_text,exclude_words=exclude_words,k=k)
   sim_dict_finbert = calculate_top_k(outputs_finbert, new_tokens_finbert, text, kp_dict_finbert, exclude_text=exclude_text,exclude_words=exclude_words,k=k)
+  sim_dict_sapbert = calculate_top_k(outputs_sapbert, new_tokens_sapbert, text, kp_dict_sapbert, exclude_text=exclude_text,exclude_words=exclude_words,k=k)
   if not diversify_box:
     with col1:
       st.write("distilbert-cvent")
@@ -112,11 +127,16 @@ if text:
     with col2:
       st.write("finbert")
       st.json(sim_dict_finbert)
+   with col3:
+    st.write("sapbert")
+    st.json(sim_dict_sapbert)
   else:
     idxs = extract_idxs(sim_dict, kp_dict)
     idxs_finbert = extract_idxs(sim_dict_finbert, kp_dict_finbert)
+    idxs_sapbert = extract_idxs(sim_dict_sapbert, kp_dict_sapbert)
     distances_candidates = cosine_kp[np.ix_(idxs, idxs)]
     distances_candidates_finbert = cosine_finbert_kp[np.ix_(idxs_finbert, idxs_finbert)]
+    distances_candidates_sapbert = cosine_sapbert_kp[np.ix_(idxs_sapbert, idxs_sapbert)]
     #first do distilbert
     candidate = None
     min_sim = np.inf
@@ -133,6 +153,14 @@ if text:
       if sim < min_sim:
         candidate_finbert = combination
         min_sim = sim
+    #sapbert
+    candidate_sapbert = None
+    min_sim = np.inf
+    for combination in itertools.combinations(range(len(idxs_sapbert)), k):
+      sim = sum([distances_candidates_sapbert[i][j] for i in combination for j in combination if i != j])
+      if sim < min_sim:
+        candidate_sapbert = combination
+        min_sim = sim
     #distilbert
     ret = {keys[idxs[idx]]:sim_dict[keys[idxs[idx]]] for idx in candidate}
     ret = sorted(ret.items(), key= lambda x: x[1], reverse = True)
@@ -141,9 +169,16 @@ if text:
     ret_finbert = {keys_finbert[idxs_finbert[idx]]:sim_dict_finbert[keys_finbert[idxs_finbert[idx]]] for idx in candidate_finbert}
     ret_finbert = sorted(ret_finbert.items(), key= lambda x: x[1], reverse = True)
     ret_finbert = {x:y for x,y in ret_finbert}
+    #sapbert
+    ret_sapbert = {keys_sapbert[idxs_sapbert[idx]]:sim_dict_sapbert[keys_sapbert[idxs_sapbert[idx]]] for idx in candidate_sapbert}
+    ret_sapbert = sorted(ret_sapbert.items(), key= lambda x: x[1], reverse = True)
+    ret_sapbert = {x:y for x,y in ret_sapbert}
     with col1:
       st.write("distilbert-cvent")
       st.json(ret)
     with col2:
       st.write("finbert")
-      st.json(ret_finbert)
+      st.json(ret_finbert)
+    with col3:
+      st.write("sapbert")
+      st.json(ret_sapbert)