Update model_functions.py (#4)

- Update model_functions.py (92c19487aec2d0780de454861d86578e1eea9fc9)


Co-authored-by: Marcelo Moreno <[email protected]>

model_functions.py

@@ -3,16 +3,17 @@ from transformers import (AutoModelForSequenceClassification, AutoModelForSeq2Se
                           AutoConfig, AutoModelForTokenClassification,
                           AutoTokenizer, pipeline)
 from peft import PeftModel, PeftConfig
+import streamlit as st
 
 
-
-
+@st.cache_resource
 def load_sentiment_analyzer():
     tokenizer = AutoTokenizer.from_pretrained("aliciiavs/sentiment-analysis-whatsapp2")
     model = AutoModelForSequenceClassification.from_pretrained("aliciiavs/sentiment-analysis-whatsapp2")
 
     return tokenizer, model
 
+@st.cache_resource
 def load_summarizer():
     config = PeftConfig.from_pretrained("marcelomoreno26/bart-large-samsum-adapter")
     model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-large")
@@ -23,6 +24,7 @@ def load_summarizer():
 
     return tokenizer, model
 
+@st.cache_resource
 def load_NER():
     config = AutoConfig.from_pretrained("hannahisrael03/wikineural-multilingual-ner-finetuned-wikiann")
     model = AutoModelForTokenClassification.from_pretrained("hannahisrael03/wikineural-multilingual-ner-finetuned-wikiann",config=config)
@@ -96,7 +98,4 @@ def get_NER(text, pipe):
             # Add non-ORG entities without filtering
             filtered_results.append(entity)
 
-    return filtered_results
-
-
-
+    return filtered_results