change defaults etc

app.py

@@ -5,14 +5,20 @@ import plotly.express as px
 from models import NLI_MODEL_OPTIONS, NSP_MODEL_OPTIONS, METHOD_OPTIONS
 from zeroshot_turkish.classifiers import NSPZeroshotClassifier, NLIZeroshotClassifier
 
-if "current_model" not in st.session_state:
-    st.session_state["current_model"] = None
 
-if "current_model_option" not in st.session_state:
-    st.session_state["current_model_option"] = None
+def init_state(key: str):
+    if key not in st.session_state:
+        st.session_state[key] = None
 
-if "current_method_option" not in st.session_state:
-    st.session_state["current_method_option"] = None
+
+for k in [
+    "current_model",
+    "current_model_option",
+    "current_method_option",
+    "current_prediction",
+    "current_chart",
+]:
+    init_state(k)
 
 
 def load_model(model_option: str, method_option: str, random_state: int = 0):
@@ -66,13 +72,11 @@ method_option = st.radio(
 )
 if method_option == METHOD_OPTIONS["nli"]:
     model_option = st.selectbox(
-        "Select a natural language inference model.",
-        NLI_MODEL_OPTIONS,
+        "Select a natural language inference model.", NLI_MODEL_OPTIONS, index=3
     )
 if method_option == METHOD_OPTIONS["nsp"]:
     model_option = st.selectbox(
-        "Select a BERT model for next sentence prediction.",
-        NSP_MODEL_OPTIONS,
+        "Select a BERT model for next sentence prediction.", NSP_MODEL_OPTIONS, index=0
     )
 
 if model_option != st.session_state.current_model_option:
@@ -105,17 +109,25 @@ prompt_template = col2.text_area(
     key="current_template",
 )
 col2.header("")
+
+
 make_pred = col1.button("Predict")
 if make_pred:
-    prediction = st.session_state.current_model.predict_on_texts(
-        [st.session_state.current_text],
-        candidate_labels=st.session_state.current_labels.split(","),
-        prompt_template=st.session_state.current_template,
+    st.session_state.current_prediction = (
+        st.session_state.current_model.predict_on_texts(
+            [st.session_state.current_text],
+            candidate_labels=st.session_state.current_labels.split(","),
+            prompt_template=st.session_state.current_template,
+        )
     )
-    if "scores" in prediction[0]:
-        chart = visualize_output(prediction[0]["labels"], prediction[0]["scores"])
-    elif "probabilities" in prediction[0]:
-        chart = visualize_output(
-            prediction[0]["labels"], prediction[0]["probabilities"]
+    if "scores" in st.session_state.current_prediction[0]:
+        st.session_state.current_chart = visualize_output(
+            st.session_state.current_prediction[0]["labels"],
+            st.session_state.current_prediction[0]["scores"],
+        )
+    elif "probabilities" in st.session_state.current_prediction[0]:
+        st.session_state.current_chart = visualize_output(
+            st.session_state.current_prediction[0]["labels"],
+            st.session_state.current_prediction[0]["probabilities"],
         )
-    col2.plotly_chart(chart)
+    col2.plotly_chart(st.session_state.current_chart, use_container_width=True)