import streamlit as st
from transformers import pipeline

# Load the HuggingFace pipeline with a binary classification model
# Use a pipeline as a high-level helper
models = {
    "fBert Convabuse": "alexabrahall/fbert_convabuse_ss",
    "fBert HTDM": "alexabrahall/fbert_htdm_ss",
    "hateBert Convabuse": "alexabrahall/hatebert_convabuse_ss",
    "hateBert HTDM": "alexabrahall/hatebert_htdm_ss",
    "berTweet Convabuse": "alexabrahall/bertweet_convabuse_ss",
    "berTweet HTDM": "alexabrahall/bertweet_htdm_ss",
    "roberta Convabuse": "alexabrahall/roberta_convabuse_ss",
    "roberta HTDM": "alexabrahall/roberta_htdm_ss",
    
}

model_descriptions = {
    "fBert Convabuse": "This is the model fBert, trained on the conversational abuse public dataset. It is a binary classification model that predicts whether a given text is abusive or not. The model is based on the fBert architecture and was trained using the Sentence Transformers library.",
    "fBert HTDM": "This is the model fBert, trained on the hate speech public dataset. It is a binary classification model that predicts whether a given text is hate speech or not. The model is based on the fBert architecture and was trained using the Sentence Transformers library.",
    "hateBert Convabuse": "This is the model hateBert, trained on the conversational abuse public dataset. It is a binary classification model that predicts whether a given text is abusive or not. The model is based on the hateBert architecture and was trained using the Sentence Transformers library.",
    "hateBert HTDM": "This is the model hateBert, trained on the hate speech public dataset. It is a binary classification model that predicts whether a given text is hate speech or not. The model is based on the hateBert architecture and was trained using the Sentence Transformers library.",
    "berTweet Convabuse": "This is the model berTweet, trained on the conversational abuse public dataset. It is a binary classification model that predicts whether a given text is abusive or not. The model is based on the berTweet architecture and was trained using the Sentence Transformers library.",
    "berTweet HTDM": "This is the model berTweet, trained on the hate speech public dataset. It is a binary classification model that predicts whether a given text is hate speech or not. The model is based on the berTweet architecture and was trained using the Sentence Transformers library.",
    "roberta Convabuse": "This is the model roberta, trained on the conversational abuse public dataset. It is a binary classification model that predicts whether a given text is abusive or not. The model is based on the roberta architecture and was trained using the Sentence Transformers library.",
    "roberta HTDM": "This is the model roberta, trained on the hate speech public dataset. It is a binary classification model that predicts whether a given text is hate speech or not. The model is based on the roberta architecture and was trained using the Sentence Transformers library.",
}


def classify_text(text, model):
    classifier = pipeline("text-classification", model=model)
    # Use the model to predict the class of the input text
    results = classifier(text)
    return results

# Streamlit app layout
st.title("Homotransphobia Detection App")
st.markdown("[By Alex Abrahall](https://huggingface.co/alexabrahall)", unsafe_allow_html=True)
st.write("Enter the text you want to classify:")

# Input box for user text
user_input = st.text_input("")


selected_model = st.selectbox("Select Model", options=list(models.keys()))
selected_model_description = model_descriptions[selected_model]
st.write("Model Description:")
st.write(selected_model_description)

label_map ={
    "LABEL_0": "Not Homotransphobic",
    "LABEL_1": "Homotransphobic"

}


# Button to classify text
if st.button('Classify'):
    prediction_raw_text = st.empty()
    prediction_text = st.empty()
    
    
    loading_text = st.text("Predicting... (if the model has not been used before, this may take a while)")
    # Classify the text
    prediction = classify_text(user_input, models[selected_model])

    loading_text.empty()
    
    # Display the result
    prediction_raw_text.write("Prediction:")
    print(prediction)
    prediction_text.markdown(f"The text is <span style='color: {'red' if prediction[0]['label'] == 'LABEL_1' else 'green'}'>{label_map[prediction[0]['label']]}</span> with a confidence of {prediction[0]['score']*100}%", unsafe_allow_html=True)