Update app.py

app.py

@@ -7,18 +7,30 @@ import pickle
 import numpy as np
 import itertools
 
-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")
+@st.cache
+def load_bert():
+  return (AutoModelForMaskedLM.from_pretrained("vives/distilbert-base-uncased-finetuned-cvent-2019_2022", output_hidden_states=True),
+  AutoTokenizer.from_pretrained("vives/distilbert-base-uncased-finetuned-cvent-2019_2022"))
+
+model, tokenizer = load_bert()
 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")
+
+@st.cache
+def load_finbert():
+  return (AutoModelForSequenceClassification.from_pretrained("ProsusAI/finbert", output_hidden_states=True),
+  AutoTokenizer.from_pretrained("ProsusAI/finbert"))
+  
+model_finbert, tokenizer_finbert = load_finbert()
 kp_dict_finbert_checkpoint = "kp_dict_finance.pickle"
 kp_cosine_finbert_checkpoint = "cosine_kp_finance.pickle"
 
-tokenizer_sapbert = AutoTokenizer.from_pretrained("cambridgeltl/SapBERT-from-PubMedBERT-fulltext")
-model_sapbert = AutoModel.from_pretrained("cambridgeltl/SapBERT-from-PubMedBERT-fulltext", output_hidden_states=True)
+@st.cache
+def load_sapbert():
+  return (AutoModel.from_pretrained("cambridgeltl/SapBERT-from-PubMedBERT-fulltext", output_hidden_states=True),
+  AutoTokenizer.from_pretrained("cambridgeltl/SapBERT-from-PubMedBERT-fulltext"))
+  
+model_sapbert, tokenizer_sapbert = load_sapbert()
 kp_dict_sapbert_checkpoint = "kp_dict_medical.pickle"
 kp_cosine_sapbert_checkpoint = "cosine_kp_medical.pickle"