Upload distilbert.py

distilbert.py

@@ -0,0 +1,117 @@
+# -*- coding: utf-8 -*-
+"""DistilBERT.ipynb
+
+Automatically generated by Colab.
+
+Original file is located at
+    https://colab.research.google.com/drive/1qXwFT-lCqgfmQYxeJ7cb-iuvTLqLkiim
+"""
+
+#DISTILLBERT RUN 3 , added weight_decay=0.01
+import pandas as pd
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import classification_report
+from transformers import BertTokenizer
+
+
+# Load dataset
+file_path = 'spam_ham_dataset.csv'
+df = pd.read_csv(file_path)
+
+# Convert labels to numeric
+df['label_num'] = df['label'].map({'ham': 0, 'spam': 1})
+
+# Load tokenizer
+tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
+
+# Tokenize dataset
+encodings = tokenizer(df['text'].tolist(), padding=True, truncation=True, max_length=128, return_tensors="pt")
+labels = torch.tensor(df['label_num'].values)
+
+# Custom Dataset
+class SpamDataset(Dataset):
+    def __init__(self, encodings, labels):
+        self.encodings = encodings
+        self.labels = labels
+
+    def __len__(self):
+        return len(self.labels)
+
+    def __getitem__(self, idx):
+        item = {key: val[idx] for key, val in self.encodings.items()}
+        item['labels'] = torch.tensor(self.labels[idx], dtype=torch.long)
+        return item
+
+# Create dataset
+dataset = SpamDataset(encodings, labels)
+
+# Split dataset (80% train, 20% validation)
+train_size = int(0.8 * len(dataset))
+val_size = len(dataset) - train_size
+train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])
+
+# DataLoader with batch size
+def collate_fn(batch):
+    keys = batch[0].keys()
+    return {key: torch.stack([b[key] for b in batch]) for key in keys}
+
+train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True, collate_fn=collate_fn)
+val_loader = DataLoader(val_dataset, batch_size=16, shuffle=False, collate_fn=collate_fn)
+
+# Load DistilBERT model
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=2)
+model.to(device)
+
+# Define optimizer and loss function
+optimizer = optim.AdamW(model.parameters(), lr=5e-5, weight_decay=0.01)
+loss_fn = nn.CrossEntropyLoss()
+
+# Training Loop
+EPOCHS = 10
+for epoch in range(EPOCHS):
+    model.train()
+    total_loss = 0
+
+    for batch in train_loader:
+        optimizer.zero_grad()
+
+        inputs = {key: val.to(device) for key, val in batch.items()}
+        labels = inputs.pop("labels").to(device)
+
+        outputs = model(**inputs)
+        loss = loss_fn(outputs.logits, labels)
+
+        loss.backward()
+        optimizer.step()
+
+        total_loss += loss.item()
+
+    avg_loss = total_loss / len(train_loader)
+    print(f"Epoch {epoch+1}, Loss: {avg_loss:.4f}")
+
+# Save trained model
+torch.save(model.state_dict(), "distilbert_spam_model.pt")
+
+# Evaluation
+model.eval()
+correct = 0
+total = 0
+with torch.no_grad():
+    for batch in val_loader:
+        inputs = {key: val.to(device) for key, val in batch.items()}
+        labels = inputs.pop("labels").to(device)
+
+        outputs = model(**inputs)
+        predictions = torch.argmax(outputs.logits, dim=1)
+        correct += (predictions == labels).sum().item()
+        total += labels.size(0)
+
+accuracy = correct / total
+print(f"Validation Accuracy: {accuracy:.4f}")