Adding plotly plot remove plot if loading

app.py

@@ -1,35 +1,20 @@
-"""
-HuggingFace Spaces that:
-- loads in HanmunRoBERTa model https://huggingface.co/bdsl/HanmunRoBERTa
-- optionally strips text of punctuation and unwanted charactesr
-- predicts century for the input text
-- Visualizes prediction scores for each century
-
-# https://huggingface.co/blog/streamlit-spaces
-# https://huggingface.co/docs/hub/en/spaces-sdks-streamlit
-
-"""
-
 import streamlit as st
 from transformers import pipeline
 from string import punctuation
 import pandas as pd
 import plotly.express as px
 import plotly.graph_objects as go
-colors = px.colors.qualitative.Plotly
 
-# from huggingface_hub import InferenceClient
-# client = InferenceClient(model="bdsl/HanmunRoBERTa")
+# Initialize or retrieve the session state variable
+if 'plot_visible' not in st.session_state:
+    st.session_state.plot_visible = False  # Initially, the plot is not visible
 
 def strip_input_str(x):
-    # Specify the characters to remove
     characters_to_remove = "○□()〔〕:\"。·, ?ㆍ" + punctuation
     translating = str.maketrans('', '', characters_to_remove)
     x = x.translate(translating)
-
     return x.strip()
 
-
 # Load the pipeline with the HanmunRoBERTa model
 model_pipeline = pipeline(task="text-classification", model="bdsl/HanmunRoBERTa")
 
@@ -47,42 +32,42 @@ input_str = st.text_area(
     height=150, 
     value="權知 高麗 國事臣某言。 伏惟小邦, 自 恭愍王 無嗣薨逝之後, 辛旽 子 禑 冒姓竊位者.")
 
-# Remove punctuation if checkbox is selected
 if remove_punct and input_str:
     input_str = strip_input_str(input_str)
-
-    # Display the input text after processing
     st.write("Processed input:", input_str)
 
-# Predict and display the classification scores if input is provided
+# Button to classify the text and toggle the visibility of the plot
 if st.button("Classify"):
+    st.session_state.plot_visible = not st.session_state.plot_visible  # Toggle the plot visibility
     if input_str:
-        predictions = model_pipeline(input_str, top_k=None)
-        data = pd.DataFrame(predictions)
-        data=data.sort_values(by='score', ascending=True)
-        data.label = data.label.astype(str)
-
-        
-        # Displaying predictions as a bar chart
-        fig = go.Figure(
-            go.Bar(
-                x=data.score.values,
-                y=[f'{i}th Century' for i in data.label.values],
-                orientation='h', 
-                text=[f'{score:.3f}' for score in data['score'].values],  # Format text with 2 decimal points
-                textposition='outside',  # Position the text outside the bars
-                hoverinfo='text',  # Use custom text for hover info
-                hovertext=[f'{i}th Century<br>Score: {score:.3f}' for i, score in zip(data['label'], data['score'])],  # Custom hover text
-                marker=dict(color=[colors[i % len(colors)] for i in range(len(data))]),  # Cycle through colors
-
-            ))
-
-        fig.update_layout(
-            height=300,  # Custom height
-            xaxis_title='Score',
-            yaxis_title='',
-            title='Model predictions and scores',
-            uniformtext_minsize=8, 
-            uniformtext_mode='hide',
-        )
-        st.plotly_chart(figure_or_data=fig, use_container_width=True)
+        with st.spinner("Classifying..."):
+            predictions = model_pipeline(input_str, top_k=None)
+            data = pd.DataFrame(predictions)
+            data = data.sort_values(by='score', ascending=True)
+            data.label = data.label.astype(str)
+
+        # Ensure the plot is only displayed when `plot_visible` is True
+        if st.session_state.plot_visible:
+            colors = px.colors.qualitative.Plotly
+            fig = go.Figure(
+                go.Bar(
+                    x=data.score.values,
+                    y=[f'{i}th Century' for i in data.label.values],
+                    orientation='h',
+                    text=[f'{score:.3f}' for score in data['score'].values],
+                    textposition='outside',
+                    hoverinfo='text',
+                    hovertext=[f'{i}th Century<br>Score: {score:.3f}' for i, score in zip(data['label'], data['score'])],
+                    marker=dict(color=[colors[i % len(colors)] for i in range(len(data))]),
+                ))
+
+            fig.update_layout(
+                height=300,
+                xaxis_title='Score',
+                yaxis_title='',
+                title='Model predictions and scores',
+                uniformtext_minsize=8, 
+                uniformtext_mode='hide',
+            )
+            st.plotly_chart(figure_or_data=fig, use_container_width=True)
+            st.session_state.plot_visible = False  # Reset to False after displaying