import streamlit as st
import inference
import pandas as pd
import const

st.title("Japanese to Emotions")
st.write(
    "I fine-tuned the BERT-based distillation model for classification of Japanese text."
)

if "input_text" not in st.session_state:
    st.session_state.input_text = ""

input_text = st.text_area(
    "Japanese text", value=st.session_state.input_text, max_chars=512
)

suggestions = ["今日は愛犬と散歩した", "猫カフェ行きたい", "自転車盗まれた"]

COLUMNS_NUM = len(suggestions)
cols = st.columns(COLUMNS_NUM)

for i, suggestion in enumerate(suggestions):
    with cols[i]:
        if st.button(suggestion, use_container_width=True):
            st.session_state.input_text = suggestion
            st.rerun()

if input_text:
    probs_dict = inference.exec(input_text)

    label_dict = {
        const.EMOTIONS[0]: "😊 Joy",
        const.EMOTIONS[1]: "😢 Sadness",
        const.EMOTIONS[2]: "😮 Anticipation",
        const.EMOTIONS[3]: "😲 Surprise",
        const.EMOTIONS[4]: "😠 Anger",
        const.EMOTIONS[5]: "😨 Fear",
        const.EMOTIONS[6]: "😖 Disgust",
        const.EMOTIONS[7]: "😉 Trust",
    }

    df = pd.DataFrame(
        {
            "Emotion": label_dict.values(),
            "Probs": [probs_dict[emotion] for emotion in const.EMOTIONS],
        }
    )

    st.bar_chart(df.set_index("Emotion"), horizontal=True)

st.write(
    """
- [Model](https://huggingface.co/koshin2001/Japanese-to-emotions)
- [Blog](https://zenn.dev/koshin/articles/6b27acdf8bbe01)
- [GitHub](https://github.com/koshin01/emotion_classification)
"""
)