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


def app():
    st.title('Toxic Comment Detector')
    st.write('Модель: ruBert tiny toxicity.')
    
    st.image('https://media4.giphy.com/media/CdhxVrdRN4YFi/giphy.gif')

    model_checkpoint = 'cointegrated/rubert-tiny-toxicity'
    tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
    model = AutoModelForSequenceClassification.from_pretrained(model_checkpoint)
    
    if torch.cuda.is_available():
        model.cuda()
    
    def text2toxicity(text, aggregate=True):
        """ Calculate toxicity of a text (if aggregate=True) or a vector of toxicity aspects (if aggregate=False)"""
        with torch.no_grad():
            inputs = tokenizer(text, return_tensors='pt', truncation=True, padding=True).to(model.device)
            proba = torch.sigmoid(model(**inputs).logits).cpu().numpy()
        if isinstance(text, str):
            proba = proba[0]
        if aggregate:
            return 1 - proba.T[0] * (1 - proba.T[-1])
        return proba
    
    user_input = st.text_area("Enter your text:", "Собака сутулая")
    if st.button("Analyze"):
        toxicity_score = text2toxicity(user_input, True)
        st.write(f"Toxicity Score: {toxicity_score:.4f}")
        if toxicity_score > 0.5:
            st.write("Warning: The text seems to be toxic!")