app.py

# app.py
import streamlit as st
from transformers import pipeline
from llama import load_llama_model, generate_llama_summary, PROMPT_TEMPLATE

st.set_page_config(page_title="Email Summarizer", layout="wide")

# Predefined examples with reference summaries
EXAMPLES = {
    "Sample1": {
        "email": """Sehr geehrte Damen und Herren,
hiermit möchte ich dass all meine personenbezogenen Daten (Vertragskonto 401384807) gelöscht werden.
Ich bitte um eine schriftliche Bestätigung darüber, dass alle Daten von mir mit dem o.g. Vertragskonto gelöscht worden sind.
Mit freundlichen Grüßen,
Liselotte Metz
""",
        "reference": "Der Kunde fordert die Löschung aller personenbezogenen Daten und eine schriftliche Bestätigung dazu."
    },
    "Sample2": {
        "email": """
        Kunde: 	Wernecke GmbH & Co. OHG und Niemeier Carsten Stiftung & Co. KG
Verbrauchsstelle:	Klappstr. 33
Zähler:			DE71117317620819513570
Vertragsnummer: 408591713
Sehr geehrte Damen und Herren,
wir bedanken uns für Ihre Abschlagsanpassung gemäß Ihres Schreibens vom 29.07.2023, 
auch wenn Ihr Schreiben auf unsere eingelegten Widersprüche nicht eingegangen ist.
Der guten Ordnung halber verweisen wir erneut zu Ihrer im Dezember 2022 angekündigten Preiserhöhung auf unseren Widerspruch vom 6. Januar 2023, den wir nach wie vor aus dargelegtem Grund aufrecht erhalten. 
Mit freundlichen Grüßen
Dr. Abbas Fechner
- Hausverwaltung -
""",
        "reference": "Der Kunde verweist auf den weiterhin bestehenden Widerspruch gegen die Preiserhöhung vom Dezember 2022 und erwartet eine Prüfung des Widerspruchs und Stellungnahme."
    }
}

@st.cache_resource
def load_all_models():
    """Pre-load all models during app initialization"""
    with st.spinner("Loading models... This may take a few minutes"):
        models = {
            "mt5-small": pipeline(
                "summarization", 
                model="ak2603/mt5-small-synthetic-data-plus-translated"
            ),
            "Llama 3.2": load_llama_model(),
            "Llama 7b Instruct": None  # Placeholder
        }
    return models

# Initialize models when app loads
models = load_all_models()

# Streamlit UI Configuration
st.title("📧 Automated Email Summarization")

# Sidebar Controls
with st.sidebar:
    st.header("Configuration")
    model_choice = st.selectbox(
        "Select Model",
        ["mt5-small", "Llama 3.2", "Llama 7b Instruct"],
        index=0
    )
    
    st.markdown("---")
    st.markdown("**Model Information:**")
    st.info(f"Selected model: {model_choice}")
    st.info(f"Total loaded models: {len([m for m in models.values() if m is not None])}")

# Main Content Area
col1, col2 = st.columns([2, 1])

with col1:
    st.subheader("Input Email")
    
    # Example selection
    example_choice = st.selectbox(
        "Choose an example email:",
        ["Custom Email"] + list(EXAMPLES.keys()),
        index=0
    )
    
    # Set email text based on selection
    email_text = EXAMPLES[example_choice]["email"] if example_choice in EXAMPLES else ""
    
    email_input = st.text_area(
        "Paste your email here:",
        value=email_text,
        height=300,
        key="input_text",
        placeholder="Enter email content here..."
    )

with col2:
    st.subheader("Summary Generation")
    if st.button("Generate Summary", use_container_width=True):
        if not email_input:
            st.error("Please enter some email content first!")
        else:
            try:
                selected_model = models[model_choice]
                
                if selected_model is None:
                    st.error("Selected model is not implemented yet")
                else:
                    with st.spinner("Generating summary..."):
                        if model_choice == "mt5-small":
                            result = selected_model(
                                email_input,
                                max_length=150,
                                do_sample=True,
                                repetition_penalty=1.5
                            )[0]['summary_text']
                        elif model_choice == "Llama 3.2":
                            model_obj, tokenizer = selected_model
                            result = generate_llama_summary(
                                email_input,
                                model_obj,
                                tokenizer,
                                PROMPT_TEMPLATE
                            )
                        
                        # Display results
                        st.success("**Generated Summary:**")
                        st.write(result)
                        
                        # Show reference summary if example is selected
                        if example_choice in EXAMPLES:
                            st.success("**Reference Summary:**")
                            st.write(EXAMPLES[example_choice]["reference"])
                        
                        # Add export options
                        st.download_button(
                            label="Download Summary",
                            data=result,
                            file_name="email_summary.txt",
                            mime="text/plain"
                        )
                        
            except Exception as e:
                st.error(f"Error generating summary: {str(e)}")

# Footer
st.markdown("---")
st.markdown("_Automated email summarization system v1.0_")