Update app.py

app.py

@@ -16,6 +16,42 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 # 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()}  # Keep as PyTorch tensors
+        item['labels'] = torch.tensor(self.labels[idx], dtype=torch.long)  # Ensure labels are `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 Function (Fix Collate)
+def collate_fn(batch):
+    keys = batch[0].keys()
+    collated = {key: torch.stack([b[key] for b in batch]) for key in keys}
+    return collated
+
+# Create DataLoader
+train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True, collate_fn=collate_fn)
+val_loader = DataLoader(val_dataset, batch_size=8, shuffle=False, collate_fn=collate_fn)
+
 # Load the trained model
 def load_model(model_path="distilbert_spam_model.pt"):
     model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=2)