Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
| import torch | |
| model_name = "Rockinrumble/spam_classifier_bert" | |
| def load_model(): | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| model = AutoModelForSequenceClassification.from_pretrained(model_name) | |
| return tokenizer,model | |
| def predict(text): | |
| inputs = tokenizer(text,return_tensors = "pt",truncation = True,padding = True) | |
| outputs = model(**inputs) | |
| probs = torch.nn.functional.softmax(outputs,logits,dim = 1) | |
| pred = torch.argmax(probs , dim = 1).item() | |
| return "Spam" if pred == 1 else "Not Spam" | |
| st.title("Spam Detector - BERT") | |
| st.write("Paste an email message to check if it's spam") | |
| user_input = st.text_area("Email Content", height = 200) | |
| if st.button("Classify"): | |
| if not user_input.strip(): | |
| st.warning("Please enter some text") | |
| else: | |
| result = predict(user_input) | |
| st.sucess(f"Prediction: {result}") |