import streamlit as st
from transformers import BertForSequenceClassification, BertTokenizer, TextClassificationPipeline

model_path = "JungleLee/bert-toxic-comment-classification"
tokenizer = BertTokenizer.from_pretrained(model_path)
model = BertForSequenceClassification.from_pretrained(model_path, num_labels=2)

pipeline = TextClassificationPipeline(model=model, tokenizer=tokenizer, return_all_scores=False)

st.set_page_config(page_title="Message Vibe Checker", page_icon="😬")
st.title("🕵️‍♂️ Check Your Message Before Sending!")
st.subheader("Don't let your message get you canceled 😅")

user_message = st.text_area("Paste your message here:", placeholder="e.g., You're a freaking genius!")

if st.button("Analyze Message"):
    if not user_message.strip():
        st.warning("You forgot to type something, buddy!")
    else:
        with st.spinner("Scanning your message for rage and sass..."):
            result = pipeline(user_message)[0]
            label = result['label']
            score = result['score']

        if label == "toxic":
            st.error(f"☢️ Whoa there! That message might be **toxic** (confidence: {score:.2f})")
            st.caption("Try sending good vibes instead 🧘‍♂️🌈")
        else:
            st.success(f"✅ Looks good! This message seems **clean** (confidence: {score:.2f})")
            st.balloons()
            st.caption("Spread kindness like confetti 🎉")

st.markdown("---")
st.caption("Made with love")