import streamlit as st
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
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
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 is 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}")