Hugging Face's logo

Spaces:
Pau22
/
Toxic_Comment_Classifier_using_DistilBERT
Build error

App Files Community
Toxic_Comment_Classifier_using_DistilBERT
File size: 3,222 Bytes
f658b05
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
import streamlit as st
from transformers import pipeline
import re
import pandas as pd

# ----------------------------
# Load model only once
# ----------------------------
@st.cache_resource
def load_model():
    model_id = "Pau22/distilbert-toxic-model"
    return pipeline(
        "text-classification",
        model=model_id,
        tokenizer=model_id,
        top_k=1  # prevents nested lists
    )

classifier = load_model()

# ----------------------------
# Clean input text
# ----------------------------
def clean_text(text):
    text = re.sub(r"http\S+|www\S+", "", str(text))
    text = text.encode("ascii", "ignore").decode()  # remove emojis
    return re.sub(r"\s+", " ", text).strip()

# ----------------------------
# Label Mapping
# ----------------------------
LABEL_MAP = {"LABEL_0": "Not Toxic", "LABEL_1": "Toxic"}

# ----------------------------
# Streamlit UI
# ----------------------------
st.set_page_config(page_title="Toxic Comment Classifier", layout="centered")
st.title("🧠 Toxic Comment Classifier — DistilBERT by Pau22")
st.write("Detects whether a comment is **Toxic** or **Not Toxic** using a fine-tuned DistilBERT model.")

# Example text options
toxic_samples = [
    "You are the worst person ever.",
    "Shut up you idiot.",
    "You f__king clown.",
    "Nobody likes you, go away.",
]

non_toxic_samples = [
    "Have a lovely day!",
    "Thank you for your help!",
    "I appreciate your effort.",
    "This was very helpful, thanks!",
]

col1, col2 = st.columns(2)

with col1:
    toxic_choice = st.selectbox("Choose a Toxic Example (Optional)", ["-- None --"] + toxic_samples)

with col2:
    non_toxic_choice = st.selectbox("Choose a Non-Toxic Example (Optional)", ["-- None --"] + non_toxic_samples)

user_text = ""
if toxic_choice != "-- None --":
    user_text = toxic_choice
elif non_toxic_choice != "-- None --":
    user_text = non_toxic_choice

user_text = st.text_area("Enter your text for analysis", user_text, height=120)

# ----------------------------
# Prediction button
# ----------------------------
if st.button("Predict"):
    if user_text.strip() == "":
        st.warning("Please enter or select a comment.")
    else:
        cleaned = clean_text(user_text)

        raw = classifier(cleaned)

        # Normalize HuggingFace output (fixes the TypeError)
        if isinstance(raw, list):
            if len(raw) > 0 and isinstance(raw[0], list):
                raw = raw[0][0]
            else:
                raw = raw[0]

        label = LABEL_MAP.get(raw["label"], raw["label"])
        score = float(raw["score"])

        st.subheader("Prediction")
        st.markdown(f"### **{label}**")
        st.write(f"Confidence: **{score:.3f}**")
        st.progress(score)

        with st.expander("Raw Model Output"):
            st.json(raw)

# ----------------------------
# Model Metrics Section
# ----------------------------
st.markdown("---")
st.subheader("📊 Model Evaluation")

metrics = {
    "Metric": ["Loss", "Accuracy", "Precision", "Recall", "F1 Score"],
    "Value": [0.1062, 0.9685, 0.8337, 0.8292, 0.8314],
}

df = pd.DataFrame(metrics)
st.table(df)

st.caption("Model trained for 2 epochs on the Jigsaw Toxic Comment Dataset.")