ayushsinha commited on
Commit
0dfc28e
·
verified ·
1 Parent(s): 1d68149

Upload 2 files

Browse files
Files changed (2) hide show
  1. app.py +54 -0
  2. requirements.txt +4 -0
app.py ADDED
@@ -0,0 +1,54 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import BertTokenizer, BertForSequenceClassification
3
+ import torch
4
+
5
+ # Load the tokenizer and model
6
+ model_name = "AventIQ-AI/bert-spam-detection"
7
+ tokenizer = BertTokenizer.from_pretrained(model_name)
8
+ model = BertForSequenceClassification.from_pretrained(model_name)
9
+
10
+ # Set the model to evaluation mode and move it to the appropriate device
11
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
12
+ model.to(device)
13
+ model.eval()
14
+
15
+ # Define the prediction function
16
+ def predict_spam(text):
17
+ """Predicts whether a given text is spam or not."""
18
+ # Tokenize input text
19
+ inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
20
+ inputs = {key: value.to(device) for key, value in inputs.items()}
21
+
22
+ # Perform inference
23
+ with torch.no_grad():
24
+ outputs = model(**inputs)
25
+ logits = outputs.logits
26
+ probabilities = torch.softmax(logits, dim=1)
27
+ prediction = torch.argmax(probabilities, dim=1).item()
28
+ confidence = probabilities[0][prediction].item()
29
+
30
+ # Map prediction to label
31
+ label_map = {0: "Not Spam", 1: "Spam"}
32
+ result = f"Prediction: {label_map[prediction]}\nConfidence: {confidence:.2f}"
33
+ return result
34
+
35
+ # Create the Gradio interface
36
+ iface = gr.Interface(
37
+ fn=predict_spam,
38
+ inputs=gr.Textbox(label="📧 Input Text", placeholder="Enter the email or message content here...", lines=5),
39
+ outputs=gr.Textbox(label="🔍 Spam Detection Result"),
40
+ title="🛡️ BERT-Based Spam Detector",
41
+ description="Enter the content of an email or message to determine whether it's Spam or Not Spam.",
42
+ examples=[
43
+ ["Congratulations! You've won a $1,000,000 lottery. Click here to claim your prize."],
44
+ ["Hey, are we still meeting for lunch tomorrow?"],
45
+ ["URGENT: Your account has been compromised. Reset your password immediately by clicking this link."],
46
+ ["Don't miss out on our exclusive offer! Buy one, get one free on all items."],
47
+ ["Can you send me the report by end of the day? Thanks!"]
48
+ ],
49
+ theme="compact",
50
+ allow_flagging="never"
51
+ )
52
+
53
+ if __name__ == "__main__":
54
+ iface.launch()
requirements.txt ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ torch
2
+ transformers
3
+ gradio
4
+ sentencepiece