testing new alternative to caching and state and models and tokenizers

src/main.py

@@ -1,15 +1,10 @@
 import streamlit as st
 from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-# Title
-st.title("CSGY-6613 Sentiment Analysis")
-# Subtitle
-st.markdown("### Ryan Kim (rk2546)")
-st.markdown("")
-
-@st.cache(allow_output_mutation=True)
-def load_model(model_name):
-    return pipeline(model=model_name, task="sentiment-analysis")
+# We'll be using Torch this time around
+import torch
+import torch.nn.functional as F
 
 @st.cache(allow_output_mutation=True)
 def label_dictionary(model_name):
@@ -24,14 +19,38 @@ def label_dictionary(model_name):
         return twitter_roberta
     return lambda x: x
 
+@st.cache(allow_output_mutation=True)
+def load_model(model_name):
+    model = AutoModelForSequenceClassification.from_pretrained(model_name)
+    tokenizer = AutoTokenizer.from_pretrained("cardiffnlp/twitter-roberta-base-sentiment")
+    classifier = pipeline(task="sentiment-analysis", model=model, tokenizer=tokenizer)
+    parser = label_dictionary(model_name)
+    return model, tokenizer, classifier, parser
+
+# We first initialize a state. The state will include the following:
+# 1) the name of the model (default: cardiffnlp/twitter-roberta-base-sentiment)
+# 2) the model itself, and 
+# 3) the parser for the outputs, in case we actually need to parse the output to something more sensible
 if "model" not in st.session_state:
     st.session_state.model_name = "cardiffnlp/twitter-roberta-base-sentiment"
-    st.session_state.model = load_model("cardiffnlp/twitter-roberta-base-sentiment")
-    st.session_state.label_parser = label_dictionary("cardiffnlp/twitter-roberta-base-sentiment")
+    model, tokenizer, classifier, parser = load_model("cardiffnlp/twitter-roberta-base-sentiment")
+    st.session_state.model = model
+    st.session_state.tokenizer = tokenizer
+    st.session_state.classifier = classifier
+    st.session_state.parser = parser
 
 def model_change():
-    st.session_state.model = load_model(st.session_state.model_name)
-    st.session_state.label_parser = label_dictionary(st.session_state.model_name)
+    model, tokenizer, classifier, parser = load_model(st.session_state.model_name)
+    st.session_state.model = model
+    st.session_state.tokenizer = tokenizer
+    st.session_state.classifier = classifier
+    st.session_state.parser = parser
+
+# Title
+st.title("CSGY-6613 Sentiment Analysis")
+# Subtitle
+st.markdown("### Ryan Kim (rk2546)")
+st.markdown("")
 
 model_option = st.selectbox(
     "What sentiment analysis model do you want to use?",
@@ -58,11 +77,11 @@ if submit:
     st.markdown("> {}".format(to_eval))
     st.write("Using the NLP model:")
     st.markdown("> {}".format(st.session_state.model_name))
-    result = st.session_state.model(to_eval)  
+    result = st.session_state.classifier(to_eval)  
     label = result[0]['label']
     score = result[0]['score']
 
-    label = st.session_state.label_parser(label)
+    label = st.session_state.parser(label)
 
     st.markdown("#### Result:")
     st.markdown("**{}**: {}".format(label,score))