Update app.py

app.py

@@ -7,9 +7,16 @@ import pickle
 import numpy as np
 import itertools
 
-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)
+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)
+else if choice == "finbert":
+  model_checkpoint = "ProsusAI/finbert"
+  tokenizer = AutoTokenizer.from_pretrained("ProsusAI/finbert")
+  model = AutoModelForSequenceClassification.from_pretrained("ProsusAI/finbert", output_hidden_states=True)
+
 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')")