Update app.py

app.py

@@ -55,9 +55,12 @@ with open(kp_cosine_finbert_checkpoint,'rb') as 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):
+def calculate_top_k(out, tokens,text,kp_dict,exclude_text=False,exclude_words=False, k=5, pooler=True):
   sim_dict = {}
-  pools = pool_embeddings(out, tokens).detach().numpy()
+  if pooler:
+    pools = pool_embeddings(out, tokens).detach().numpy()
+  else:
+    pools = out["pooler_outputs"].detach().numpy()
   for key in kp_dict.keys():
     if key == text:
       continue
@@ -116,7 +119,7 @@ if text:
   with torch.no_grad():
     outputs = model(**new_tokens)
     outputs_finbert = model_finbert(**new_tokens_finbert)
-    outputs_sapbert = model_sapbert(**new_tokens_sapbert)
+    outputs_sapbert = model_sapbert(**new_tokens_sapbert, pooler=False)
   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)