trying to make programmatic way to store sentiment analysis label dictionaries for each type

src/main.py

@@ -11,12 +11,27 @@ st.markdown("")
 def load_model(model_name):
     return pipeline(model=model_name, task="sentiment-analysis")
 
+@st.cache(allow_output_mutation=True)
+def label_dictionary(model_name):
+    if model_name == "cardiffnlp/twitter-roberta-base-sentiment":
+        def twitter_roberta(label):
+            if label == "LABEL_0":
+                return "Negative"
+            elif label == "LABEL_2":
+                return "Positive"
+            else:
+                return "Neutral"
+        return twitter_roberta
+    return lambda x: x
+
 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")
 
 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_option = st.selectbox(
     "What sentiment analysis model do you want to use?",
@@ -44,12 +59,7 @@ if submit:
     label = result[0]['label']
     score = result[0]['score']
 
-    if label == "LABEL_0":
-        label = "Negative"
-    elif label == "LABEL_2":
-        label = "Positive"
-    elif label == "LABEL_1":
-        label = "Neutral"
+    label = st.session_state.label_parser(label)
 
     st.markdown("#### Result:")
     st.markdown("**{}**: {}".format(label,score))