Adding plotly plot

app.py

@@ -1,33 +1,83 @@
-import gradio as gr
+"""
+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
 
-title = "RoBERTa"
+# https://huggingface.co/blog/streamlit-spaces
+# https://huggingface.co/docs/hub/en/spaces-sdks-streamlit
 
-description = "Gradio Demo for RoBERTa. To use it, simply add your text, or click one of the examples to load them. Read more at the links below."
+"""
 
-article = "<p style='text-align: center'><a href='https://arxiv.org/abs/1907.11692' target='_blank'>RoBERTa: A Robustly Optimized BERT Pretraining Approach</a></p>"
+import streamlit as st
+from transformers import pipeline
+from string import punctuation
+import pandas as pd
+import plotly.express as px
+import plotly.graph_objects as go
+colors = px.colors.qualitative.Plotly
 
-examples = [
-    ['The goal of life is <mask>.','roberta-base']
-]
+# from huggingface_hub import InferenceClient
+# client = InferenceClient(model="bdsl/HanmunRoBERTa")
 
-io1 = gr.Interface.load("huggingface/roberta-base")
+# Load the pipeline with the HanmunRoBERTa model
+model_pipeline = pipeline(task="text-classification", model="bdsl/HanmunRoBERTa")
 
-io2 = gr.Interface.load("huggingface/roberta-large")
+# Streamlit app layout
+title = "HanmunRoBERTa Century Classifier"
+st.set_page_config(page_title=title, page_icon="📚")
+st.title(title)
 
+# Checkbox to remove punctuation
+remove_punct = st.checkbox(label="Remove punctuation", value=True)
 
-def inference(inputtext, model):
-    if model == "roberta-base":
-        outlabel = io1(inputtext)
+# 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)
+        data = pd.DataFrame(predictions)
+        data=data.sort_values(by='score', ascending=True)
+        data.label = data.label.astype(str)
+
+        
+        # Displaying predictions as a bar chart
+        fig = go.Figure(
+            go.Bar(
+                x=data.score.values,
+                y=[f'{i}th Century' for i in data.label.values],
+                orientation='h', 
+                text=[f'{score:.3f}' for score in data['score'].values],  # Format text with 2 decimal points
+                textposition='outside',  # Position the text outside the bars
+                hoverinfo='text',  # Use custom text for hover info
+                hovertext=[f'{i}th Century<br>Score: {score:.3f}' for i, score in zip(data['label'], data['score'])],  # Custom hover text
+                marker=dict(color=[colors[i % len(colors)] for i in range(len(data))]),  # Cycle through colors
+
+            ))
+        fig.update_traces(width=0.4)
+
+        fig.update_layout(
+            height=300,  # Custom height
+            xaxis_title='Score',
+            yaxis_title='',
+            title='Model predictions and scores',
+            margin=dict(l=100, r=200, t=50, b=50),
+            uniformtext_minsize=8, 
+            uniformtext_mode='hide',
+        )
+        st.pyplot(fig=fig)
     else:
-        outlabel = io2(inputtext)
-    return outlabel
-     
-
-gr.Interface(
-    inference, 
-    [gr.Textbox(label="Context",lines=10),gr.Dropdown(choices=["roberta-base","roberta-large"], type="value",  label="model")], 
-    [gr.Label(label="Output")],
-    examples=examples,
-    article=article,
-    title=title,
-    description=description).launch(enable_queue=True)
+        st.write("Please enter some text to classify.")

requirements.txt

@@ -1,4 +1,5 @@
 streamlit
 torch
 transformers
-pandas
+pandas
+plotly