Create app.py

app.py

@@ -0,0 +1,63 @@
+import pandas as pd
+import streamlit as st
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
+from datasets import load_dataset
+
+# Load the model and tokenizer
+model_name = "modelSamLowe/roberta-base-go_emotions"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSequenceClassification.from_pretrained(model_name)
+
+# Define the emotion labels (based on the GoEmotions dataset)
+emotion_labels = ["admiration", "amusement", "anger", "annoyance", "approval",
+                  "caring", "confusion", "curiosity", "desire", "disappointment",
+                  "disapproval", "disgust", "embarrassment", "excitement", "fear",
+                  "gratitude", "grief", "joy", "love", "nervousness", "optimism",
+                  "pride", "realization", "relief", "remorse", "sadness", "surprise",
+                  "neutral"]
+
+# Function to classify emotion
+def classify_emotion(text):
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+    outputs = model(**inputs)
+    logits = outputs.logits
+    predicted_class_id = torch.argmax(logits, dim=-1).item()
+    return emotion_labels[predicted_class_id]
+
+# Streamlit interface
+st.title("Enron Emails Emotion Analysis")
+
+# Button to run the inference script
+if st.button("Run Inference"):
+    # Load the Enron dataset
+    with st.spinner('Loading dataset...'):
+        dataset = load_dataset("Hellisotherpeople/enron_emails_parsed")
+        enron_data = pd.DataFrame(dataset['train'])
+
+    # Apply emotion classification to the email content
+    with st.spinner('Running inference...'):
+        enron_data['emotion'] = enron_data['body'].apply(classify_emotion)
+
+    # Save the results to a CSV file
+    enron_data.to_csv("enron_emails_with_emotions.csv", index=False)
+    st.success("Inference completed and results saved!")
+
+# Check if the results file exists and load it
+try:
+    enron_data = pd.read_csv("enron_emails_with_emotions.csv")
+
+    # Dropdown for selecting an emotion
+    selected_emotion = st.selectbox("Select Emotion", emotion_labels)
+
+    # Filter emails based on the selected emotion
+    filtered_emails = enron_data[enron_data['emotion'] == selected_emotion].head(10)
+
+    # Display the filtered emails in a table
+    if not filtered_emails.empty:
+        st.write("Top 10 emails with emotion:", selected_emotion)
+        st.table(filtered_emails[['From', 'To', 'body', 'emotion']])
+    else:
+        st.write("No emails found with the selected emotion.")
+except FileNotFoundError:
+    st.warning("Run inference first by clicking the 'Run Inference' button.")