Update app.py

app.py

@@ -56,14 +56,22 @@ class NERLabelEncoder:
 NER_CHECKPOINT = "microsoft/deberta-base"
 NER_N_TOKENS = 50
 NER_N_LABELS = 18
-ner_model = TFAutoModelForTokenClassification.from_pretrained(NER_CHECKPOINT, num_labels=NER_N_LABELS, attention_probs_dropout_prob=0.4, hidden_dropout_prob=0.4)
-ner_model.load_weights(os.path.join("models", "general_ner_deberta_weights.h5"), by_name=True)
-ner_label_encoder = NERLabelEncoder()
-ner_label_encoder.fit()
-ner_tokenizer = DebertaTokenizerFast.from_pretrained(NER_CHECKPOINT, add_prefix_space=True)
-nlp = spacy.load(os.path.join('.', 'en_core_web_sm-3.6.0'))
 NER_COLOR_MAP = {'GEO': '#DFFF00', 'GPE': '#FFBF00', 'PER': '#9FE2BF', 
                  'ORG': '#40E0D0', 'TIM': '#CCCCFF', 'ART': '#FFC0CB', 'NAT': '#FFE4B5', 'EVE': '#DCDCDC'}
+
+@st.cache_resource
+def load_ner_models():
+    ner_model = TFAutoModelForTokenClassification.from_pretrained(NER_CHECKPOINT, num_labels=NER_N_LABELS, attention_probs_dropout_prob=0.4, hidden_dropout_prob=0.4)
+    ner_model.load_weights(os.path.join("models", "general_ner_deberta_weights.h5"), by_name=True)
+    ner_label_encoder = NERLabelEncoder()
+    ner_label_encoder.fit()
+    ner_tokenizer = DebertaTokenizerFast.from_pretrained(NER_CHECKPOINT, add_prefix_space=True)
+    nlp = spacy.load(os.path.join('.', 'en_core_web_sm-3.6.0'))
+    print('Loaded NER models')
+    return ner_model, ner_label_encoder, ner_tokenizer, nlp
+
+ner_model, ner_label_encoder, ner_tokenizer, nlp = load_ner_models()
+
 ############ NER MODEL & VARS INITIALIZATION END ####################
 
 ############ NER LOGIC START ####################
@@ -170,9 +178,16 @@ def get_ner_text(article_txt, ner_result):
 SUMM_CHECKPOINT = "facebook/bart-base"
 SUMM_INPUT_N_TOKENS = 400
 SUMM_TARGET_N_TOKENS = 100
-summ_tokenizer = BartTokenizerFast.from_pretrained(SUMM_CHECKPOINT)
-summ_model = TFAutoModelForSeq2SeqLM.from_pretrained(SUMM_CHECKPOINT)
-summ_model.load_weights(os.path.join("models", "bart_en_summarizer.h5"), by_name=True)
+
+@st.cache_resource
+def load_summarizer_models():
+    summ_tokenizer = BartTokenizerFast.from_pretrained(SUMM_CHECKPOINT)
+    summ_model = TFAutoModelForSeq2SeqLM.from_pretrained(SUMM_CHECKPOINT)
+    summ_model.load_weights(os.path.join("models", "bart_en_summarizer.h5"), by_name=True)
+    print('Loaded summarizer models')
+    return summ_tokenizer, summ_model
+
+summ_tokenizer, summ_model = load_summarizer_models()
 
 def summ_preprocess(txt):
     txt = re.sub(r'^By \. [\w\s]+ \. ', ' ', txt) # By . Ellie Zolfagharifard . 
@@ -190,7 +205,6 @@ def summ_preprocess(txt):
     return txt
 
 def summ_inference_tokenize(input_: list, n_tokens: int):
-    # tokenizer = BartTokenizerFast.from_pretrained(SUMM_CHECKPOINT)
     tokenized_data = summ_tokenizer(text=input_, max_length=SUMM_TARGET_N_TOKENS, truncation=True, padding="max_length", return_tensors="tf")
     return summ_tokenizer, tokenized_data    
 
@@ -207,7 +221,7 @@ def summ_inference(txt: str):
 ############## ENTRY POINT START #######################
 def main():
     st.title("News Summarizer & NER")
-    article_txt = st.text_area("Paste the text of a news article:", "", height=200)
+    article_txt = st.text_area("Paste few sentences of a news article:", "", height=200)
     if st.button("Submit"):
         ner_result = [[ent, label.upper(), np.round(prob, 3)] 
                                   for ent, label, prob in ner_inference_long_text(article_txt)]