import streamlit as st
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

# Load the model and tokenizer
model_name = "Hate-speech-CNERG/urdu-abusive-MuRIL"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Define labels in Urdu
labels = {0: "نارمل (معمول)", 1: "گالی گلوچ (بدتمیزی)"}

# Streamlit UI
st.title("اردو بدتمیزی کا تجزیہ")
st.write("یہ ایپ آپ کے متن کا تجزیہ کرے گی اور بتائے گی کہ یہ نارمل ہے یا بدتمیز۔")

# Input text from user
user_input = st.text_area("اپنا متن یہاں درج کریں:")

if st.button("تجزیہ کریں"):
    if user_input.strip():
        # Tokenize and classify the input
        inputs = tokenizer(user_input, return_tensors="pt", truncation=True, padding=True)
        outputs = model(**inputs)
        logits = outputs.logits
        predicted_class = torch.argmax(logits, dim=1).item()

        # Display the result
        st.write(f"متن کی نوعیت: **{labels[predicted_class]}**")
    else:
        st.warning("براہ کرم متن درج کریں!")