Update app.py

app.py

@@ -1,13 +1,19 @@
 import pandas as pd
-import streamlit as st
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
 from datasets import load_dataset
 
+# Check if GPU is available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Load the Enron dataset
+dataset = load_dataset("Hellisotherpeople/enron_emails_parsed")
+enron_data = pd.DataFrame(dataset['train'])
+
 # Load the model and tokenizer
-model_name = "SamLowe/roberta-base-go_emotions"
+model_name = "modelSamLowe/roberta-base-go_emotions"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
+model = AutoModelForSequenceClassification.from_pretrained(model_name).to(device)
 
 # Define the emotion labels (based on the GoEmotions dataset)
 emotion_labels = ["admiration", "amusement", "anger", "annoyance", "approval",
@@ -17,47 +23,22 @@ emotion_labels = ["admiration", "amusement", "anger", "annoyance", "approval",
                   "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.")
+# Function to classify emotions in batches
+def classify_emotions_in_batches(texts, batch_size=32):
+    results = []
+    for i in range(0, len(texts), batch_size):
+        batch = texts[i:i+batch_size]
+        inputs = tokenizer(batch, return_tensors="pt", truncation=True, padding=True).to(device)
+        with torch.no_grad():
+            outputs = model(**inputs)
+            logits = outputs.logits
+            predicted_class_ids = torch.argmax(logits, dim=-1).tolist()
+            results.extend(predicted_class_ids)
+    return results
+
+# Apply emotion classification to the email content in batches
+email_texts = enron_data['body'].tolist()
+enron_data['emotion'] = classify_emotions_in_batches(email_texts, batch_size=32)
+
+# Save the results to a CSV file
+enron_data.to_csv("enron_emails_with_emotions.csv", index=False)