update app

app.py

@@ -5,7 +5,6 @@ import pickle
 import pandas as pd
 from transformers import RobertaTokenizerFast, RobertaModel
 
-# -------------------------------
 # Load label mappings
 with open("label_mappings.pkl", "rb") as f:
     label_mappings = pickle.load(f)
@@ -13,11 +12,10 @@ with open("label_mappings.pkl", "rb") as f:
 label_to_team = label_mappings.get("label_to_team", {})
 label_to_email = label_mappings.get("label_to_email", {})
 
-# -------------------------------
+
 # Load tokenizer
 tokenizer = RobertaTokenizerFast.from_pretrained("roberta-base")
 
-# -------------------------------
 # Define RoBERTa Model for multi-task classification
 class RoBertaClassifier(nn.Module):
     def __init__(self, num_teams, num_emails):
@@ -33,7 +31,6 @@ class RoBertaClassifier(nn.Module):
         email_logits = self.email_classifier(cls_output)
         return team_logits, email_logits
 
-# -------------------------------
 # Initialize model and load checkpoint
 num_teams = len(label_to_team)
 num_emails = len(label_to_email)
@@ -47,7 +44,7 @@ device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cp
 model.to(device)
 model.eval()
 
-# -------------------------------
+
 # Prediction function
 def predict_tickets(ticket_descriptions):
     predictions = []
@@ -73,19 +70,23 @@ def predict_tickets(ticket_descriptions):
     df = pd.DataFrame(csv_data, columns=["Index", "Description", "Assigned Team", "Team Email"])
     return "\n".join(predictions), df
 
-# -------------------------------
+
 # Streamlit UI
-#st.title("AI Solution for Defect Ticket Classification")
-st.markdown("<h2 style='text-align: center;'>AI Solution for Defect Ticket Classification</h2>", unsafe_allow_html=True)
+st.markdown("<h2 style='text-align: center; font-size:22px;'>AI Solution for Defect Ticket Classification</h2>", unsafe_allow_html=True)
+
 st.markdown("""
-**Supports:** Multi-line text input & CSV upload.  
-**Output:** Text results & downloadable CSV file.  
-**Model:** Fine-tuned **RoBERTa** for classification.
-""")
+<p style='text-align: center; font-size:16px;'><strong>Supports:</strong> Multi-line text input & CSV upload.</p>  
+<p style='text-align: center; font-size:16px;'><strong>Output:</strong> Text results & downloadable CSV file.</p>  
+<p style='text-align: center; font-size:16px;'><strong>Model:</strong> Fine-tuned <strong>RoBERTa</strong> for classification.</p>
+""", unsafe_allow_html=True)
+
+st.markdown("<h3 style='font-size:16px;'>Enter ticket Description/Comment/Summary or upload a CSV file to predict Assigned Team & Team Email.</h3>", unsafe_allow_html=True)
+
 
 # Choose input method
 option = st.radio("📝 Choose Input Method", ["Enter Text", "Upload CSV"])
 
+descriptions = []
 if option == "Enter Text":
     text_input = st.text_area(
         "Enter Ticket Description/Comment/Summary (One per line)", 
@@ -94,7 +95,6 @@ if option == "Enter Text":
     descriptions = [line.strip() for line in text_input.split("\n") if line.strip()]
 else:
     file_input = st.file_uploader("Upload CSV", type=["csv"])
-    descriptions = []
     if file_input is not None:
         df_input = pd.read_csv(file_input)
         if "Description" not in df_input.columns:
@@ -102,29 +102,47 @@ else:
         else:
             descriptions = df_input["Description"].dropna().tolist()
 
-# Trigger prediction when the button is clicked
-if st.button("PREDICT"):
-    if not descriptions:
-        st.error("⚠️ Please provide valid input.")
-    else:
-        with st.spinner("Predicting..."):
-            results, df_results = predict_tickets(descriptions)
-        st.markdown("## Prediction Results")
-        st.text(results)
-        csv_data = df_results.to_csv(index=False).encode('utf-8')
-        st.download_button(
-            label="📥 Download Predictions CSV",
-            data=csv_data,
-            file_name="ticket-predictions.csv",
-            mime="text/csv"
-        )
 
-# Clear button: simply reloads the app
-if st.button("CLEAR"):
-    st.experimental_rerun()
+# Store prediction results in session state so they persist
+if "prediction_results" not in st.session_state:
+    st.session_state.prediction_results = None
+if "df_results" not in st.session_state:
+    st.session_state.df_results = None
+
+# Create a horizontal layout for the buttons
+col1, col2 = st.columns([1, 1])
+
+with col1:
+    if st.button("PREDICT"):
+        if not descriptions:
+            st.error("⚠️ Please provide valid input.")
+        else:
+            with st.spinner("Predicting..."):
+                results, df_results = predict_tickets(descriptions)
+            st.session_state.prediction_results = results
+            st.session_state.df_results = df_results
+
+# Display prediction results if available
+if st.session_state.prediction_results:
+    st.markdown("<h3 style='font-size:16px;'>Prediction Results</h3>", unsafe_allow_html=True)
+    st.text(st.session_state.prediction_results)
+    csv_data = st.session_state.df_results.to_csv(index=False).encode('utf-8')
+    st.download_button(
+        label="📥 Download Predictions CSV",
+        data=csv_data,
+        file_name="ticket-predictions.csv",
+        mime="text/csv"
+    )
+
+with col2:
+    if st.button("CLEAR"):
+        # Clear the prediction results from session state
+        st.session_state.prediction_results = None
+        st.session_state.df_results = None
+        st.rerun()
 
 st.markdown("---")
 st.markdown(
-    "<p style='text-align: center;color: gray;'>Developed by NYP student @ Min Thein Win: Student ID: 3907578Y</p>", 
+    "<p style='text-align: center;color: gray; font-size:14px;'>Developed by NYP student @ Min Thein Win: Student ID: 3907578Y</p>", 
     unsafe_allow_html=True
-)
+)