import streamlit as st
from bias_detector import Detector

st.title("Multidimensional Stereotype Multilevel Detection")

EXAMPLES = {
    "Gender": "The baby loved the presence of his caring mommy.",
    "Religion": "The orange church had a very good reputation.",
    "Race": "Steve thinks singaporeans are short, but he doesn't know any better",
    "Profession": "The football player's draft card said he was tough.",
    "Input your own": ""
    # Add more examples as needed
}
def format_results(results):
    formatted = ""
    for result in results:
        for text, pred in result.items():
            formatted += f"**Text**: {text}\n\n"
            formatted += "**Predictions**:\n"
            # Check if pred is a dictionary
            if isinstance(pred, dict):
                for token, labels in pred.items():
                    if isinstance(labels, dict):  # handling token-level output
                        formatted += f"- Token: `{token}`\n"
                        for label, score in labels.items():
                            formatted += f"    - Label: `{label}`, Score: `{score}`\n"
                    else:  # handling sentence-level output when labels is not a dictionary
                        formatted += f"- Label: `{token}`, Score: `{labels}`\n"
            else:  # handling edge cases where pred is not a dictionary
                formatted += f"Prediction score: {pred}\n"
    return formatted



level = st.selectbox("Select the Bias Levels:", ("Sentence","Token"))
dimension = st.selectbox("Select the Bias Dimensions:", ("All","Gender","Religion","Race","Profession"))

detector = Detector(level,dimension)

if st.button("Load Models"):
    st.text("Loading models...")

    dummy_sentence = "This is a dummy sentence."
    dummy_result = detector.predict([dummy_sentence])  # perform a dummy prediction
    if dummy_result:
        st.text("Models loaded successfully!")
    else:
        st.text("Failed to load models. Please check the server and/or model parameters.")

example_type = st.selectbox("Choose an example type:", list(EXAMPLES.keys()))
target_sentence = st.text_input("Input the sentence you want to detect:", value=EXAMPLES[example_type])

if st.button("Detect"):
    results = detector.predict([target_sentence])
    if results:
        formatted_results = format_results(results)
        st.markdown(f"## Detection Results: \n\n {formatted_results}")
    else:
        st.text("Prediction failed. Please check the input and try again.")