Update app.py

app.py

@@ -21,10 +21,14 @@ emotion_labels = ["admiration", "amusement", "anger", "annoyance", "approval",
                   "pride", "realization", "relief", "remorse", "sadness", "surprise",
                   "neutral"]
 
-# Function to classify emotions in batches
-def classify_emotions_in_batches(texts, batch_size=32):
+# Function to classify emotions in batches and save results after each batch
+def classify_emotions_in_batches(texts, batch_size=64, output_file="enron_emails_with_emotions.csv"):
     results = []
     start_time = time.time()
+
+    # DataFrame to store the results
+    result_df = pd.DataFrame(columns=['From', 'To', 'body', 'emotion'])
+
     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)
@@ -33,12 +37,26 @@ def classify_emotions_in_batches(texts, batch_size=32):
             logits = outputs.logits
             predicted_class_ids = torch.argmax(logits, dim=-1).tolist()
             results.extend(predicted_class_ids)
-        
+
+        # Save the batch results
+        batch_results = {
+            'From': enron_data['From'][i:i+batch_size],
+            'To': enron_data['To'][i:i+batch_size],
+            'body': batch,
+            'emotion': [emotion_labels[idx] for idx in predicted_class_ids]
+        }
+        batch_df = pd.DataFrame(batch_results)
+        result_df = pd.concat([result_df, batch_df])
+
+        # Save to CSV
+        result_df.to_csv(output_file, index=False)
+
         # Log progress
         batch_time = time.time() - start_time
         st.write(f"Processed batch {i//batch_size + 1} of {len(texts)//batch_size + 1} in {batch_time:.2f} seconds")
         start_time = time.time()
-    return results
+
+    return result_df
 
 # Streamlit interface
 st.title("Enron Emails Emotion Analysis")
@@ -53,28 +71,6 @@ if st.button("Run Inference"):
     # Apply emotion classification to the email content
     with st.spinner('Running inference...'):
         email_texts = enron_data['body'].tolist()
-        enron_data['emotion'] = classify_emotions_in_batches(email_texts, batch_size=32)
+        classify_emotions_in_batches(email_texts, batch_size=64)
 
-    # 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.")
-