leygit/ITI110_Spam_Classification_Project

1# -*- coding: utf-8 -*-2"""DistilBERT.ipynb3 4Automatically generated by Colab.5 6Original file is located at7    https://colab.research.google.com/drive/1qXwFT-lCqgfmQYxeJ7cb-iuvTLqLkiim8"""9 10#DISTILLBERT RUN 3 , added weight_decay=0.0111import pandas as pd12import torch13import torch.nn as nn14import torch.optim as optim15import torch.nn.functional as F16from torch.utils.data import Dataset, DataLoader17from transformers import DistilBertTokenizer, DistilBertForSequenceClassification18from sklearn.model_selection import train_test_split19from sklearn.metrics import classification_report20from transformers import BertTokenizer21 22 23# Load dataset24file_path = 'spam_ham_dataset.csv'25df = pd.read_csv(file_path)26 27# Convert labels to numeric28df['label_num'] = df['label'].map({'ham': 0, 'spam': 1})29 30# Load tokenizer31tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')32 33# Tokenize dataset34encodings = tokenizer(df['text'].tolist(), padding=True, truncation=True, max_length=128, return_tensors="pt")35labels = torch.tensor(df['label_num'].values)36 37# Custom Dataset38class SpamDataset(Dataset):39    def __init__(self, encodings, labels):40        self.encodings = encodings41        self.labels = labels42 43    def __len__(self):44        return len(self.labels)45 46    def __getitem__(self, idx):47        item = {key: val[idx] for key, val in self.encodings.items()}48        item['labels'] = torch.tensor(self.labels[idx], dtype=torch.long)49        return item50 51# Create dataset52dataset = SpamDataset(encodings, labels)53 54# Split dataset (80% train, 20% validation)55train_size = int(0.8 * len(dataset))56val_size = len(dataset) - train_size57train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])58 59# DataLoader with batch size60def collate_fn(batch):61    keys = batch[0].keys()62    return {key: torch.stack([b[key] for b in batch]) for key in keys}63 64train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True, collate_fn=collate_fn)65val_loader = DataLoader(val_dataset, batch_size=16, shuffle=False, collate_fn=collate_fn)66 67# Load DistilBERT model68device = torch.device("cuda" if torch.cuda.is_available() else "cpu")69model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=2)70model.to(device)71 72# Define optimizer and loss function73optimizer = optim.AdamW(model.parameters(), lr=5e-5, weight_decay=0.01)74loss_fn = nn.CrossEntropyLoss()75 76# Training Loop77EPOCHS = 1078for epoch in range(EPOCHS):79    model.train()80    total_loss = 081 82    for batch in train_loader:83        optimizer.zero_grad()84 85        inputs = {key: val.to(device) for key, val in batch.items()}86        labels = inputs.pop("labels").to(device)87 88        outputs = model(**inputs)89        loss = loss_fn(outputs.logits, labels)90 91        loss.backward()92        optimizer.step()93 94        total_loss += loss.item()95 96    avg_loss = total_loss / len(train_loader)97    print(f"Epoch {epoch+1}, Loss: {avg_loss:.4f}")98 99# Save trained model100torch.save(model.state_dict(), "distilbert_spam_model.pt")101 102# Evaluation103model.eval()104correct = 0105total = 0106with torch.no_grad():107    for batch in val_loader:108        inputs = {key: val.to(device) for key, val in batch.items()}109        labels = inputs.pop("labels").to(device)110 111        outputs = model(**inputs)112        predictions = torch.argmax(outputs.logits, dim=1)113        correct += (predictions == labels).sum().item()114        total += labels.size(0)115 116accuracy = correct / total117print(f"Validation Accuracy: {accuracy:.4f}")