File size: 1,514 Bytes
708ef37 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
import pandas as pd
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
from datasets import load_dataset
# Load the Enron dataset
dataset = load_dataset("Hellisotherpeople/enron_emails_parsed")
enron_data = pd.DataFrame(dataset['train'])
# 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]
# Apply emotion classification to the email content
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) |