import streamlit as st
import pandas as pd
from datasets import load_dataset
from transformers import pipeline

# Load Enron Dataset
@st.cache
def load_data():

    dataset = load_dataset("Hellisotherpeople/enron_emails_parsed")
    return dataset['train']

data = load_data()

# Load models
sentiment_model = pipeline('sentiment-analysis')
ner_model = pipeline('ner', aggregation_strategy="simple")
topic_model = pipeline('zero-shot-classification', model='facebook/bart-large-mnli')

# Streamlit UI
st.title('Enron Email Analysis')

st.sidebar.title('Options')
email_id = st.sidebar.selectbox('Select Email ID', range(len(data)))

email_text = data[email_id]['body']
st.write(f"## Email Content\n{email_text}")

# Sentiment Analysis
st.write("## Sentiment Analysis")
sentiment = sentiment_model(email_text)
st.write(sentiment)

# Named Entity Recognition
st.write("## Named Entity Recognition (NER)")
entities = ner_model(email_text)
st.write(entities)

# Topic Modeling (Zero-Shot Classification)
st.write("## Topic Modeling")
labels = ['business', 'personal', 'financial', 'legal', 'politics']
topics = topic_model(email_text, candidate_labels=labels)
st.write(topics)