Update streamlit_app.py

streamlit_app.py

@@ -1,28 +1,41 @@
-# streamlit_app.py
-import streamlit as st
-import requests
-
-st.title("📧 Email Classification and PII Masking System")
-
-email_input = st.text_area("Enter your email text:")
-
-if st.button("Classify Email"):
-    if email_input.strip() != "":
-        # Send to FastAPI
-        response = requests.post("http://127.0.0.1:8000/predict", json={"email_body": email_input})
-
-        if response.status_code == 200:
-            result = response.json()
-
-            st.subheader("🔎 Masked Email:")
-            st.write(result['masked_email'])
-
-            st.subheader("🔐 List of Masked Entities:")
-            st.json(result['list_of_masked_entities'])
-
-            st.subheader("📂 Predicted Category:")
-            st.success(result['category_of_the_email'])
-        else:
-            st.error("Error from API.")
-    else:
-        st.warning("Please enter email text!")
+# streamlit_app.py
+import streamlit as st
+from models import load_model
+from utils import mask_pii
+
+st.set_page_config(page_title="📧 Email Classification System", page_icon="📧")
+
+st.title("📧 Email Classification and PII Masking System")
+
+# Load model
+model = load_model("model/classifier.pkl")
+
+# Prediction function
+def classify_email_direct(email_body):
+    masked_text, entities = mask_pii(email_body)
+    category = model.predict([masked_text])[0]
+    return {
+        "input_email_body": email_body,
+        "list_of_masked_entities": entities,
+        "masked_email": masked_text,
+        "category_of_the_email": category
+    }
+
+# Streamlit UI
+email_input = st.text_area("✉️ Enter your email text:")
+
+if st.button("🚀 Classify Email"):
+    if email_input.strip() != "":
+        result = classify_email_direct(email_input)
+
+        st.subheader("🔎 Masked Email")
+        st.write(result['masked_email'])
+
+        st.subheader("🔐 List of Masked Entities")
+        st.json(result['list_of_masked_entities'])
+
+        st.subheader("📂 Predicted Category")
+        st.success(result['category_of_the_email'])
+    else:
+        st.warning("⚠️ Please enter some text first.")
+