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

tokenizer, model = load_model()

def predict(text):
    inputs = tokenizer(text, return_tensort="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 and check if it's spam.")

user_input = st.text_area("Email Content", height=200)

if st.button("Clasify"):
    if not user_input.strip():
        st.warning("Please enter some words")
    else:
        result = predict(user_input)
        st.success(f"Prediction: {result}")