Update app.py

app.py

@@ -133,10 +133,10 @@ if text:
     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, pooler=False)
   if not diversify_box:
+    sim_dict = calculate_top_k(outputs, new_tokens, text, kp_dict, exclude_text=exclude_text,exclude_words=exclude_words,k=k_diversify)
+    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_diversify)
+    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_diversify, pooler=False)
     with col1:
       st.write("distilbert-cvent")
       st.json(sim_dict)
@@ -147,6 +147,9 @@ if text:
       st.write("sapbert")
       st.json(sim_dict_sapbert)
   else:
+    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, pooler=False)
     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)