Spaces:
Running
Running
import gradio as gr | |
from transformers import BertTokenizer, BertForSequenceClassification | |
import torch | |
# Load the tokenizer and model | |
model_name = "AventIQ-AI/bert-spam-detection" | |
tokenizer = BertTokenizer.from_pretrained(model_name) | |
model = BertForSequenceClassification.from_pretrained(model_name) | |
# Set the model to evaluation mode and move it to the appropriate device | |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
model.to(device) | |
model.eval() | |
# Define the prediction function | |
def predict_spam(text): | |
"""Predicts whether a given text is spam or not.""" | |
# Tokenize input text | |
inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512) | |
inputs = {key: value.to(device) for key, value in inputs.items()} | |
# Perform inference | |
with torch.no_grad(): | |
outputs = model(**inputs) | |
logits = outputs.logits | |
probabilities = torch.softmax(logits, dim=1) | |
prediction = torch.argmax(probabilities, dim=1).item() | |
confidence = probabilities[0][prediction].item() | |
# Map prediction to label | |
label_map = {0: "Not Spam", 1: "Spam"} | |
result = f"Prediction: {label_map[prediction]}\nConfidence: {confidence:.2f}" | |
return result | |
# Create the Gradio interface | |
iface = gr.Interface( | |
fn=predict_spam, | |
inputs=gr.Textbox(label="π§ Input Text", placeholder="Enter the email or message content here...", lines=5), | |
outputs=gr.Textbox(label="π Spam Detection Result"), | |
title="π‘οΈ BERT-Based Spam Detector", | |
description="Enter the content of an email or message to determine whether it's Spam or Not Spam.", | |
examples=[ | |
["Congratulations! You've won a $1,000,000 lottery. Click here to claim your prize."], | |
["Hey, are we still meeting for lunch tomorrow?"], | |
["URGENT: Your account has been compromised. Reset your password immediately by clicking this link."], | |
["Don't miss out on our exclusive offer! Buy one, get one free on all items."], | |
["Can you send me the report by end of the day? Thanks!"] | |
], | |
theme="compact", | |
allow_flagging="never" | |
) | |
if __name__ == "__main__": | |
iface.launch() | |