File size: 1,130 Bytes
b61909d
 
 
 
b60e88d
02bc89f
f5834da
02bc89f
ed4e35c
02bc89f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
fa0c7f2
02bc89f
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
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}")