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

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

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

# App title and description
st.title("اردو متن کا تجزیہ کریں")
st.write("یہ ایپ آپ کے فراہم کردہ اردو متن کی نوعیت (نارمل یا گالی گلوچ) کو پہچانتی ہے۔")

# User input
user_input = st.text_area("اردو متن درج کریں:")

if st.button("تجزیہ کریں"):
    if user_input.strip():
        # Tokenize and classify the input text
        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("براہ کرم متن درج کریں!")