Update app.py

app.py

@@ -36,6 +36,14 @@ def predict_class(text):
 uploaded_files_dir = "uploaded_files"
 os.makedirs(uploaded_files_dir, exist_ok=True)
 
+# Define colors for different classes
+class_colors = {
+    0: "#1f77b4",  # Level 1
+    1: "#ff7f0e",  # Level 2
+    2: "#2ca02c",  # Level 3
+    3: "#d62728"   # Level 4
+}
+
 st.title("Paper Citation Classifier")
 
 option = st.radio("Select input type:", ("Text", "PDF"))
@@ -53,7 +61,11 @@ if option == "Text":
         predicted_class = predict_class(combined_text)
         if predicted_class is not None:
             class_labels = ["Level 1", "Level 2", "Level 3", "Level 4"]
-            st.text(f"Predicted Class: {class_labels[predicted_class]}")
+            st.text("Predicted Class:")
+            st.markdown(
+                f'<div style="background-color: {class_colors[predicted_class]}; padding: 10px; border-radius: 5px; color: white; font-weight: bold;">{class_labels[predicted_class]}</div>',
+                unsafe_allow_html=True
+            )
 
 elif option == "PDF":
     uploaded_file = st.file_uploader("Upload a PDF file", type=["pdf"])
@@ -74,4 +86,8 @@ elif option == "PDF":
             predicted_class = predict_class(file_text)
             if predicted_class is not None:
                 class_labels = ["Level 1", "Level 2", "Level 3", "Level 4"]
-                st.text(f"Predicted Class: {class_labels[predicted_class]}")
+                st.text("Predicted Class:")
+                st.markdown(
+                    f'<div style="background-color: {class_colors[predicted_class]}; padding: 10px; border-radius: 5px; color: white; font-weight: bold;">{class_labels[predicted_class]}</div>',
+                    unsafe_allow_html=True
+                )