# -*- 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}")