Update app.py

app.py

@@ -7,25 +7,18 @@ import pickle
 import numpy as np
 import itertools
 
-choice = st.radio("Choose model",["distilbert-cvent","finbert"])
-if choice == "distilbert-cvent":
-  model_checkpoint = "vives/distilbert-base-uncased-finetuned-cvent-2019_2022"
-  model = AutoModelForMaskedLM.from_pretrained(model_checkpoint, output_hidden_states=True)
-  tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
-  kp_dict_checkpoint = "kp_dict_merged.pickle"
-  kp_cosine_checkpoint = "cosine_kp.pickle"
-   
-elif choice == "finbert":
-  model_checkpoint = "ProsusAI/finbert"
-  tokenizer = AutoTokenizer.from_pretrained("ProsusAI/finbert")
-  model = AutoModelForSequenceClassification.from_pretrained("ProsusAI/finbert", output_hidden_states=True)
-  kp_dict_checkpoint = "kp_dict_finbert.pickle"
-  kp_cosine_checkpoint = "cosine_kp_finbert.pickle"
+model = AutoModelForMaskedLM.from_pretrained("vives/distilbert-base-uncased-finetuned-cvent-2019_2022", output_hidden_states=True)
+tokenizer = AutoTokenizer.from_pretrained("vives/distilbert-base-uncased-finetuned-cvent-2019_2022")
+kp_dict_checkpoint = "kp_dict_merged.pickle"
+kp_cosine_checkpoint = "cosine_kp.pickle"
+ 
+model_finbert = AutoModelForSequenceClassification.from_pretrained("ProsusAI/finbert", output_hidden_states=True)
+tokenizer_finbert = AutoTokenizer.from_pretrained("ProsusAI/finbert")
+kp_dict_finbert_checkpoint = "kp_dict_finbert.pickle"
+kp_cosine_finbert_checkpoint = "cosine_kp_finbert.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 (i.e exclude 'hot coffee' if the query is 'cold coffee')")
-
+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')")
 k = st.number_input("Top k nearest key-phrases",1,10,5)
 
@@ -34,16 +27,24 @@ with st.sidebar:
   if diversify_box:
     k_diversify = st.number_input("Set of key-phrases to diversify from",10,30,20)
 
-#load kp di
+#columns
+col1, col2 = st.columns(2)
+#load kp dicts
 with open(kp_dict_checkpoint,'rb') as handle:
   kp_dict = pickle.load(handle)
 keys = list(kp_dict.keys())
 
+with open(kp_dict_finbert_checkpoint,'rb') as handle:
+  kp_dict_finbert = pickle.load(handle)
+keys_finbert = list(kp_dict_finbert.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)
   
-def calculate_top_k(out, tokens,text,exclude_text=False,exclude_words=False, k=5):
+def calculate_top_k(out, tokens,text,kp_dict,exclude_text=False,exclude_words=False, k=5):
   sim_dict = {}
   pools = pool_embeddings(out, tokens).detach().numpy()
   for key in kp_dict.keys():
@@ -59,7 +60,7 @@ def calculate_top_k(out, tokens,text,exclude_text=False,exclude_words=False, k=5
     )[0][0]
   sims = sorted(sim_dict.items(), key= lambda x: x[1], reverse = True)[:k]
   return {x:y for x,y in sims}
-def concat_tokens(sentences):
+def concat_tokens(sentences, tokenizer):
   tokens = {'input_ids': [], 'attention_mask': [], 'KPS': []}
   for sentence in sentences:
       # encode each sentence and append to dictionary
@@ -95,25 +96,54 @@ def extract_idxs(top_dict, kp_dict):
   
 if text:
   text = text.lower()
-  new_tokens = concat_tokens([text])
+  new_tokens = concat_tokens([text], tokenizer)
   new_tokens.pop("KPS")
+  new_tokens_finbert = concat_tokens([text], tokenizer_finbert)
+  new_tokens_finbert.pop("KPS")
   with torch.no_grad():
     outputs = model(**new_tokens)
+    outputs_finbert = model_finbert(**new_tokens_finbert)
+  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)
   if not diversify_box:
-    sim_dict = calculate_top_k(outputs, new_tokens, text, exclude_text=exclude_text,exclude_words=exclude_words,k=k)
-    st.json(sim_dict)
+    with col1:
+      st.write("distilbert-cvent")
+      st.json(sim_dict)
+    with col2:
+      st.write("finbert")
+      st.json(sim_dict_finbert)
   else:
-    sim_dict = calculate_top_k(outputs, new_tokens, text, exclude_text=exclude_text,exclude_words=exclude_words,k=k_diversify)
     idxs = extract_idxs(sim_dict, kp_dict)
+    idxs_finbert = extract_idxs(sim_dict, kp_dict_finbert)
     distances_candidates = cosine_kp[np.ix_(idxs, idxs)]
-    min_sim = np.inf
+    distances_candidates_finbert = cosine_kp_finbert[np.ix_(idxs_finbert, idxs_finbert)]
+    #first do distilbert
     candidate = None
+    min_sim = np.inf
     for combination in itertools.combinations(range(len(idxs)), k):
       sim = sum([distances_candidates[i][j] for i in combination for j in combination if i != j])
       if sim < min_sim:
         candidate = combination
         min_sim = sim
+    #then do finbert
+    candidate_finbert = None
+    min_sim = np.inf
+    for combination in itertools.combinations(range(len(idxs_finbert)), k):
+      sim = sum([distances_candidates_finbert[i][j] for i in combination for j in combination if i != j])
+      if sim < min_sim:
+        candidate_finbert = 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)
     ret = {x:y for x,y in ret}
-    st.json(ret)
+    #finbert
+    ret_finbert = {keys_finbert[idxs_finbert[idx]]:sim_dict_finbert[keys_finbert[idxs[idx]]] for idx in candidate_finbert}
+    candidate_finbert = sorted(candidate_finbert.items(), key= lambda x: x[1], reverse = True)
+    candidate_finbert = {x:y for x,y in candidate_finbert}
+    with col1:
+      st.write("distilbert-cvent")
+      st.json(ret)
+    with col2:
+      st.write("finbert")
+      st.json(ret_finbert)