latest

app.py

@@ -1,22 +1,21 @@
 # app.py
-import os
-import joblib
 import torch
-from transformers import BertTokenizer, BertForSequenceClassification
+import joblib
+from transformers import BertTokenizer
 from torch.nn.functional import softmax
 
-# Load the tokenizer and model
+# Load the tokenizer
 tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
 
-# Check if CUDA is available, otherwise use CPU
+# Device configuration
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load the saved model
+# Load your saved model
 model = joblib.load('model.joblib')
 model.to(device)
 model.eval()
 
-# Class names
+# Class names corresponding to the labels
 class_names = ["JAILBREAK", "INJECTION", "PHISHING", "SAFE"]
 
 def preprocess(text):