train.py

import torch
from transformers import DistilBertTokenizerFast, DistilBertForSequenceClassification, Trainer, TrainingArguments
from datasets import Dataset
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, precision_recall_fscore_support
import pandas as pd

df = pd.read_csv("final_dataset.csv", encoding='latin1')
df = df[['Email', 'Label']].dropna()

label_map = {'NORMAL': 0, 'SPAM': 1, 'FRAUD': 2}
df['label'] = df['Label'].map(label_map)

train_df, test_df = train_test_split(df, test_size=0.2, random_state=42, stratify=df['label'])

train_ds = Dataset.from_pandas(train_df[['Email', 'label']])
test_ds = Dataset.from_pandas(test_df[['Email', 'label']])

tokenizer = DistilBertTokenizerFast.from_pretrained('distilbert-base-uncased')

def tokenize(batch):
    return tokenizer(batch['Email'], truncation=True, padding='max_length', max_length=256)

train_ds = train_ds.map(tokenize, batched=True, batch_size=1024)
test_ds = test_ds.map(tokenize, batched=True, batch_size=1024)

train_ds.set_format('torch', columns=['input_ids', 'attention_mask', 'label'])
test_ds.set_format('torch', columns=['input_ids', 'attention_mask', 'label'])

model = DistilBertForSequenceClassification.from_pretrained(
    'distilbert-base-uncased', num_labels=3
)

def compute_metrics(pred):
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)
    precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average='weighted')
    acc = accuracy_score(labels, preds)
    return {'accuracy': acc, 'f1': f1, 'precision': precision, 'recall': recall}

training_args = TrainingArguments(
    output_dir='./results',
    do_eval=True,
    eval_steps=500,
    learning_rate=2e-5,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=3,
    weight_decay=0.01,
    logging_dir='./logs'
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_ds,
    eval_dataset=test_ds,
    tokenizer=tokenizer,
    compute_metrics=compute_metrics
)

trainer.train()
results = trainer.evaluate()

model.save_pretrained("./email_model")
tokenizer.save_pretrained("./email_model")