Update app.py

app.py

@@ -12,11 +12,14 @@ 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_finbert.pickle"
+  kp_cosine_checkpoint = "cosine_kp_finbert.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_merged.pickle"
+  kp_cosine_checkpoint = "cosine_kp.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')")
@@ -29,8 +32,8 @@ with st.sidebar:
   if diversify_box:
     k_diversify = st.number_input("Set of key-phrases to diversify from",10,30,20)
 
-#load kp dict
-with open("kp_dict_merged.pickle",'rb') as handle:
+#load kp di
+with open(kp_dict_checkpoint,'rb') as handle:
   kp_dict = pickle.load(handle)
 keys = list(kp_dict.keys())
 
@@ -38,7 +41,7 @@ for key in kp_dict.keys():
   kp_dict[key] = kp_dict[key].detach().numpy()
 
 #load cosine distances of kp dict
-with open("cosine_kp.pickle",'rb') as handle:
+with open(kp_cosine_checkpoint,'rb') as handle:
   cosine_kp = pickle.load(handle)
   
 def calculate_top_k(out, tokens,text,exclude_text=False,exclude_words=False, k=5):