Spam_Detector_bot / src /streamlit_app.py
Rockinrumble's picture
Update src/streamlit_app.py
fa0c7f2 verified
import os
os.environ["HF_HOME"] = "/tmp/huggingface" # or another writable path
import streamlit as st
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
#Loading model
model_name = "laiBatool/laiba-spam-classifier-bert"
@st.cache_resource
def load_model():
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
return tokenizer, model
tokenizer, model = load_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"
#Streamlit Ui
st.title("Spam Detector -BERT")
st.write("Paste an email message and check if its 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.success(f"Prediction: {result}")