"""
HuggingFace Spaces that:
- loads in HanmunRoBERTa model https://huggingface.co/bdsl/HanmunRoBERTa
- optionally strips text of punctuation and unwanted charactesr
- predicts century for the input text
- Visualizes prediction scores for each century

# https://huggingface.co/blog/streamlit-spaces
# https://huggingface.co/docs/hub/en/spaces-sdks-streamlit

"""

import streamlit as st
from transformers import pipeline
from string import punctuation
import pandas as pd
from huggingface_hub import InferenceClient
client = InferenceClient(model="bdsl/HanmunRoBERTa")

# Load the pipeline with the HanmunRoBERTa model
model_pipeline = pipeline(task="text-classification", model="bdsl/HanmunRoBERTa")

# Streamlit app layout
title = "HanmunRoBERTa Century Classifier"
st.title(title)
st.set_page_config(layout=layout, page_title=title, page_icon="📚")

# Checkbox to remove punctuation
remove_punct = st.checkbox(label="Remove punctuation", value=True)

# Text area for user input
input_str = st.text_area("Input text", height=275)

# Remove punctuation if checkbox is selected
if remove_punct and input_str:
	# Specify the characters to remove
	characters_to_remove = "○□()〔〕:\"。·, ?ㆍ" + punctuation
	translating = str.maketrans('', '', characters_to_remove)
	input_str = input_str.translate(translating)

# Display the input text after processing
st.write("Processed input:", input_str)

# Predict and display the classification scores if input is provided
if st.button("Classify"):
    if input_str:
        predictions = model_pipeline(input_str)
        
        # Prepare the data for plotting
        labels = [prediction['label'] for prediction in predictions]
        scores = [prediction['score'] for prediction in predictions]
        data = pd.DataFrame({"Label": labels, "Score": scores})
        
        # Displaying predictions as a bar chart
        st.bar_chart(data.set_index('Label'))
    else:
        st.write("Please enter some text to classify.")