Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,6 +1,6 @@
 import streamlit as st
 import torch
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer
 import time
 
 st.set_page_config(
@@ -21,18 +21,27 @@ st.markdown("""
 
 @st.cache_resource
 def load_pipelines():
+    """Load all 3 pipelines with custom classifier"""
     with st.spinner('🔄 Loading AI models (2-3 min)...'):
-        classifier = pipeline("text-classification", model="distilbert-base-uncased")
+        # ✅ YOUR CUSTOM CLASSIFIER (replaced)
+        tok = AutoTokenizer.from_pretrained("distilbert-base-uncased")
+        classifier = pipeline(
+            "text-classification",
+            model="byinab/custom-email-classifier",
+            tokenizer=tok,
+        )
+        
         generator = pipeline("text-generation", model="Kunal7370944861/Email-Writer-AI")
         translator = pipeline("translation", model="DDDSSS/translation_en-zh")
+    
     return classifier, generator, translator
 
 # Load models safely
 try:
     classifier, generator, translator = load_pipelines()
-    st.success("✅ All 3 pipelines ready!")
+    st.success("✅ All 3 pipelines ready! (Custom Email Classifier Loaded)")
 except Exception as e:
-    st.error(f"Model loading error: {str(e)}")
+    st.error(f"❌ Model loading error: {str(e)}")
     st.stop()
 
 def classify_email(text, classifier):
@@ -65,10 +74,10 @@ def translate_reply(text, translator):
 st.markdown('<h1 class="main-header">🤖 Email Reply Assistant</h1>', unsafe_allow_html=True)
 st.markdown("**AI-powered: Classify → Generate Reply → Translate to Chinese**")
 
-# Sidebar
+# Sidebar - UPDATED with your custom classifier
 with st.sidebar:
     st.header("🔧 Pipeline Status")
-    st.success("✅ **Pipeline 1**: `distilbert-base-uncased`")
+    st.success("✅ **Pipeline 1**: `byinab/custom-email-classifier`")
     st.success("✅ **Pipeline 2**: `Kunal7370944861/Email-Writer-AI`") 
     st.success("✅ **Pipeline 3**: `DDDSSS/translation_en-zh`")
     st.markdown("---")
@@ -98,10 +107,10 @@ with col2:
     if 'processed' in st.session_state and st.session_state.processed:
         email_text = st.session_state.email
         
-        # Pipeline 1: Classification
+        # Pipeline 1: YOUR CUSTOM CLASSIFIER
         with st.container():
             st.markdown('<div class="pipeline-card">', unsafe_allow_html=True)
-            st.markdown("### 🔢 **Pipeline 1: Email Classification**")
+            st.markdown("### 🔢 **Pipeline 1: Custom Email Classifier**")
             label, score = classify_email(email_text, classifier)
             
             col_a, col_b = st.columns(2)
@@ -150,4 +159,4 @@ with col2:
         st.markdown('</div>', unsafe_allow_html=True)
 
 st.markdown("---")
-st.markdown("*Powered by Streamlit + Transformers | Deployed on Hugging Face Spaces*")
+st.markdown("*Powered by Streamlit + Transformers | Custom Email Classifier*")