Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,21 +1,57 @@
-# Install (run once)
-!pip install -q transformers accelerate torch ipywidgets
-
+import streamlit as st
+import torch
 from transformers import pipeline
-from IPython.display import display, clear_output
-import ipywidgets as widgets
-
-print("Loading WORKING 3 PIPELINES...")
-
-# ✅ WORKING PIPELINES (your other 2 models + reliable classifier)
-classifier = pipeline("text-classification", model="distilbert-base-uncased")
-generator = pipeline("text-generation", model="Kunal7370944861/Email-Writer-AI")
-translator = pipeline("translation", model="DDDSSS/translation_en-zh")
-
-print("✅ ALL PIPELINES LOADED!")
+import time
+
+# Page config
+st.set_page_config(
+    page_title="📧 Email Reply Assistant",
+    page_icon="📧",
+    layout="wide",
+    initial_sidebar_state="expanded"
+)
 
-# Helper functions (unchanged)
-def classify_email(text: str):
+# Custom CSS
+st.markdown("""
+<style>
+    .main-header { font-size: 3rem; color: #1f77b4; text-align: center; margin-bottom: 2rem; }
+    .pipeline-card { background: #f8f9fa; padding: 1.5rem; border-radius: 10px; border-left: 5px solid #1f77b4; }
+    .status-success { color: #28a745; font-weight: bold; }
+    .status-loading { color: #ffc107; font-weight: bold; }
+</style>
+""", unsafe_allow_html=True)
+
+@st.cache_resource
+def load_pipelines():
+    """Load all 3 pipelines with GPU support"""
+    with st.spinner("Loading AI models... This takes ~2 minutes"):
+        classifier = pipeline(
+            "text-classification",
+            model="distilbert-base-uncased",
+            device=0 if torch.cuda.is_available() else -1
+        )
+        generator = pipeline(
+            "text-generation",
+            model="Kunal7370944861/Email-Writer-AI",
+            device=0 if torch.cuda.is_available() else -1
+        )
+        translator = pipeline(
+            "translation",
+            model="DDDSSS/translation_en-zh",
+            device=0 if torch.cuda.is_available() else -1
+        )
+    return classifier, generator, translator
+
+# Load pipelines
+try:
+    classifier, generator, translator = load_pipelines()
+    st.success("✅ All 3 AI pipelines loaded!")
+except Exception as e:
+    st.error(f"❌ Pipeline loading failed: {str(e)}")
+    st.stop()
+
+# Helper functions (same logic)
+def classify_email(text: str, classifier):
     result = classifier(text[:512])[0]
     return result["label"], float(result["score"])
 
@@ -29,90 +65,101 @@ Customer email:
 Write a polite, concise reply template.
 Reply:"""
 
-def generate_reply(prompt: str) -> str:
-    outputs = generator(prompt, max_length=300, num_return_sequences=1, do_sample=True, temperature=0.7)
+def generate_reply(prompt: str, generator):
+    outputs = generator(
+        prompt, 
+        max_length=300,
+        num_return_sequences=1, 
+        do_sample=True, 
+        temperature=0.7,
+        pad_token_id=generator.tokenizer.eos_token_id
+    )
     full_text = outputs[0]["generated_text"]
     if "Reply:" in full_text:
         return full_text.split("Reply:", 1)[-1].strip()
     return full_text.replace(prompt, "").strip()
 
-def translate_reply(text: str) -> str:
-    if not text.strip(): return ""
+def translate_reply(text: str, translator):
+    if not text.strip(): 
+        return ""
     outputs = translator(text)
     return outputs[0]["translation_text"].strip()
 
-# UI
-print("\n" + "="*70)
-print("📧 EMAIL PROCESSOR - READY!")
-print("="*70)
-
-email_input = widgets.Textarea(
-    value="", 
-    placeholder="Paste your email here...",
-    layout=widgets.Layout(width='98%', height='120px')
-)
-
-process_btn = widgets.Button(
-    description="🚀 PROCESS EMAIL", 
-    button_style='success',
-    layout=widgets.Layout(width='220px')
-)
-
-clear_btn = widgets.Button(
-    description="🧹 CLEAR", 
-    button_style='warning',
-    layout=widgets.Layout(width='120px')
-)
-
-output_box = widgets.Output()
-
-def on_process_clicked(b):
-    with output_box:
-        clear_output(wait=True)
-        email_text = email_input.value.strip()
-        
-        if not email_text:
-            print("❌ Paste an email!")
-            return
-        
-        print("📧 EMAIL")
-        print("-"*60)
-        print(email_text)
-        print("\n"+"="*60)
-        
-        # PIPELINE 1: WORKING CLASSIFIER
-        print("🔢 PIPELINE 1: distilbert-base-uncased-finetuned-sst-2-english")
-        label, score = classify_email(email_text)
-        print(f"✅ TAG: **{label}** ({score:.1%})")
-        
-        # PIPELINE 2: YOUR MODEL
-        print("\n✉️ PIPELINE 2: Kunal7370944861/Email-Writer-AI")
-        prompt = build_prompt(email_text, label)
-        reply_en = generate_reply(prompt)
-        print("✅ ENGLISH REPLY:")
-        print("-"*60)
-        print(reply_en)
-        
-        # PIPELINE 3: YOUR MODEL
-        print("\n🇨🇳 PIPELINE 3: DDDSSS/translation_en-zh")
-        reply_zh = translate_reply(reply_en)
-        print("✅ 中文回复:")
-        print("-"*60)
-        print(reply_zh)
-        
-        print("\n🎉 DONE!")
-
-def on_clear_clicked(b):
-    email_input.value = ""
-    output_box.clear_output()
-
-process_btn.on_click(on_process_clicked)
-clear_btn.on_click(on_clear_clicked)
-
-display(widgets.VBox([
-    widgets.HTML("<h3>📨 Paste Email (Blank Input)</h3>"),
-    email_input,
-    widgets.HBox([process_btn, clear_btn]),
-    widgets.HTML("<br>"),
-    output_box
-]))
+# Main title
+st.markdown('<h1 class="main-header">🤖 Email Reply Assistant</h1>', unsafe_allow_html=True)
+st.markdown("**Classify → Generate Reply → Translate to Chinese** • Powered by 3 Transformer models")
+
+# Sidebar info
+with st.sidebar:
+    st.header("ℹ️ Pipeline Info")
+    st.markdown("""
+    **Pipeline 1**: `distilbert-base-uncased` → Email classification  
+    **Pipeline 2**: `Kunal7370944861/Email-Writer-AI` → Reply generation
+    **Pipeline 3**: `DDDSSS/translation_en-zh` → Chinese translation
+    
+    **Status**: ✅ All models loaded
+    """)
+    st.markdown("---")
+    st.info("👈 Paste email → Click Process → Get instant replies!")
+
+# Main content
+col1, col2 = st.columns([1, 2])
+
+with col1:
+    st.header("📨 Input Email")
+    
+    # Email input
+    email_text = st.text_area(
+        "Paste your email here...",
+        placeholder="Subject: Problem with order\n\nHello,\n\nMy package arrived damaged...",
+        height=200,
+        help="Paste complete email (subject + body)"
+    )
+    
+    col_btn1, col_btn2 = st.columns(2)
+    with col_btn1:
+        process_btn = st.button("🚀 PROCESS EMAIL", type="primary", use_container_width=True)
+    with col_btn2:
+        if st.button("🧹 CLEAR", use_container_width=True):
+            st.rerun()
+
+with col2:
+    if process_btn and email_text.strip():
+        with st.spinner("🔄 Processing through 3 AI pipelines..."):
+            # Pipeline 1: Classification
+            with st.container():
+                st.markdown('<div class="pipeline-card"><h3>🔢 Pipeline 1: Classification</h3>', unsafe_allow_html=True)
+                label, score = classify_email(email_text, classifier)
+                col_a, col_b = st.columns(2)
+                col_a.metric("Recommended Tag", label)
+                col_b.metric("Confidence", f"{score:.1%}")
+                st.markdown('</div>', unsafe_allow_html=True)
+            
+            # Pipeline 2: Reply Generation
+            with st.container():
+                st.markdown('<div class="pipeline-card"><h3>✉️ Pipeline 2: English Reply</h3>', unsafe_allow_html=True)
+                prompt = build_prompt(email_text, label)
+                reply_en = generate_reply(prompt, generator)
+                st.text_area("English Reply Template", reply_en, height=150, disabled=True)
+                st.markdown('</div>', unsafe_allow_html=True)
+            
+            # Pipeline 3: Translation
+            with st.container():
+                st.markdown('<div class="pipeline-card"><h3>🇨🇳 Pipeline 3: Chinese Translation</h3>', unsafe_allow_html=True)
+                reply_zh = translate_reply(reply_en, translator)
+                st.text_area("中文回复模板", reply_zh, height=150, disabled=True)
+                st.markdown('</div>', unsafe_allow_html=True)
+            
+            # Copy buttons
+            st.markdown("---")
+            col_c, col_d = st.columns(2)
+            with col_c:
+                st.download_button("📥 Download English", reply_en, "reply_en.txt")
+            with col_d:
+                st.download_button("📥 Download Chinese", reply_zh, "reply_zh.txt")
+    else:
+        st.info("👆 **Paste an email and click 'PROCESS EMAIL'** to see the magic!")
+
+# Footer
+st.markdown("---")
+st.markdown("*Built with ❤️ using Streamlit + Transformers • Deployed on Hugging Face Spaces*")