switched to gradio

.gitignore

@@ -1 +1,2 @@
-data/docs.csv
+data/docs.csv
+gradio_cached_examples

README.md

@@ -3,8 +3,8 @@ title: Persplain
 emoji: 📊
 colorFrom: gray
 colorTo: green
-sdk: streamlit
-sdk_version: 1.17.0
+sdk: gradio
+sdk_version: 3.24.1
 app_file: app.py
 pinned: false
 ---

app.py

@@ -1,109 +1,109 @@
+import gradio as gr
 from transformers import RobertaForSequenceClassification, RobertaTokenizer
-from simpletransformers.classification import MultiLabelClassificationModel
 from transformers_interpret import MultiLabelClassificationExplainer
-import streamlit as st
 import pandas as pd
-import numpy as np
+from transformers import logging
+
+logging.set_verbosity_warning()
 
 traits = ["Openness to Experience", "Conscientiousness", "Extraversion", "Agreeableness", "Neuroticism" ]
 
 short_traits = ["o", "c", "e", "a", "n"]
 
-@st.experimental_memo
-def load_explainer():
+short_to_long = {"o": "Openness to Experience", "c": "Conscientiousness", "e": "Extraversion", "a": "Agreeableness", "n": "Neuroticism" }
 
+def load_explainer():
     print("Loading model...")
-
     tokenizer = RobertaTokenizer.from_pretrained("andaqu/roBERTa-pers")
     model = RobertaForSequenceClassification.from_pretrained("andaqu/roBERTa-pers", problem_type="multi_label_classification")
-
     explainer = MultiLabelClassificationExplainer(model, tokenizer)
-
     try:
         model.to('cuda')
         if next(model.parameters()).is_cuda:
             print("Using GPU for inference!")
     except:
         print("GPU not available, using CPU instead.")
-
     print("Model loaded!")
-
     return explainer
 
-@st.experimental_memo
-def explain(text, _explainer):
-    
-    attributions = _explainer(text)
-
-    preds = {label: pred_prob.item() for pred_prob, label in zip(_explainer.pred_probs_list, _explainer.labels)}
-
-    attributions_html = {trait : attributions_to_html(attributions[trait]) for trait in attributions}
+explainer = load_explainer()
 
-    return {"preds": preds, "word_attributions_html": attributions_html }
-
-def attributions_to_html(attributions):
-  html = ""
-  for word, attr in attributions:
-      
-    if word in ["<s>", "</s>"]:
-        continue
-    
-    attr = round(attr, 2)
-    abs_attr = abs(attr)
-
-    color = "rgba(255,255,255,0)"
-    if attr > 0: color = f"rgba(0,255,0,{abs_attr})"
-    elif attr < 0: color = f"rgba(255,0,0,{abs_attr})"
-      
-    html += f'<span style="background-color: {color}" title="{str(attr)}">{word}</span> '
-  
-  return html
+def explain(text, _explainer):
+    if text is not None:
+        attributions = _explainer(text)
+        preds = {label: pred_prob.item() for pred_prob, label in zip(_explainer.pred_probs_list, _explainer.labels)}
+        attributions_html = {trait : attributions_to_html(attributions[trait], trait) for trait in attributions}
+        return {"preds": preds, "word_attributions_html": attributions_html }
+    else:
+        return None
 
+def attributions_to_html(attributions, short_trait):
+    html = f""
+    for word, attr in attributions:
+        if word in ["<s>", "</s>"]:
+            continue
+        attr = round(attr, 2)
+        abs_attr = abs(attr)
+        color = "rgba(255,255,255,0)"
+        if attr > 0: color = f"rgba(0,255,0,{abs_attr})"
+        elif attr < 0: color = f"rgba(255,0,0,{abs_attr})"
+        html += f'<span style="background-color: {color}" title="{str(attr)}">{word}</span> '
+    html += f"<br>"
+    return html
 
-if "text" in st.session_state: text = st.session_state.text
-else: st.session_state.text = ""
+def get_predictions(text):
+    explanation = explain(text, explainer)
 
-if "explanation" in st.session_state: preds = st.session_state.explanation
-else: st.session_state.explanation = {"preds": {}, "word_attributions_html": ""}
+    prediction = ["YES" if explanation["preds"][x] > 0.5 else "NO" for x in explanation["preds"]]
+    probability = [str(round(explanation["preds"][x]*100)) + "%" for x in explanation["preds"]]
 
+    result_df = pd.DataFrame(data={"Predicted Traits": prediction, "Probability": probability}, index=traits)
 
-def main():
+    def color_row(row):
+        if row['Predicted Traits'] == 'YES':
+            return ['background-color: green']*len(row)
+        else:
+            return ['background-color: red']*len(row)
 
-    st.title("Text to Personality Explainer 📊")
-    text = ""
+    # apply conditional formatting to dataframe
+    result_df = result_df.style.apply(color_row, axis=1)
 
-    explainer = load_explainer()
+    def render_html(val):
+        return val
 
-    text = st.text_area(label="Input text here...", value="I enjoy meeting people and working hard!")
+    explanation_df = pd.DataFrame(data={"Explanation": [explanation["word_attributions_html"][x] for x in short_traits]}, index=traits)
 
-    show_prediction = st.button("Predict Traits")
+    explanation_df = explanation_df.style.format({'Explanation': render_html})
 
-    st.session_state.text = text
+    return result_df, explanation_df
 
-    if show_prediction and st.session_state.text:
-            
-            explanation = explain(text, explainer)
-            st.session_state.explanation = explanation
+def text_to_personality_explainer(text):
+    result_df, explanation_df = get_predictions(text)
 
-    if len(st.session_state.explanation["preds"]) > 0:
+    return "<center>" + result_df.to_html() + "</center>", "<center>" + explanation_df.to_html() + "</center>"
 
-        st.write("## Predicted Traits")
+main = gr.Blocks()
+text_input = gr.Textbox(placeholder="Enter text here...")
+result = gr.outputs.HTML()
+explanation = gr.outputs.HTML() 
 
-        prediction = ["YES" if st.session_state.explanation["preds"][x] > 0.5 else "NO" for x in st.session_state.explanation["preds"]]
-        probability = [str(round(st.session_state.explanation["preds"][x]*100)) + "%" for x in st.session_state.explanation["preds"]]
-        
-        st.table(pd.DataFrame([prediction, probability], columns=traits, index=["Predicted Traits", "Probability"]))
+with main:
+    gr.Markdown("# Text to Personality Explainer 📊")
+    gr.Markdown("Predict personality traits from text using a RoBERTa model fine-tuned on a Big Five Personality Traits dataset.")
+    gr.Markdown("Explanations are given in the form of word attributions, where the color of the word indicates the importance of the word for the prediction. Green words increase the probability of the trait, red words decrease the probability of the trait.")
 
-        st.write("## Explanation")
-        # Show five buttons, horizontally, one for each trait
-        cols = st.columns(5)
+    gr.Examples(["I love working and talking to people!", "I am a bad person. :(", "I find it challenging to agree with my brother."], fn=text_to_personality_explainer, inputs=text_input, outputs=[result, explanation], cache_examples=False)
 
-        for i in range(5):
-            button = cols[i].button(traits[i])
-        
-            if button: st.markdown(st.session_state.explanation["word_attributions_html"][short_traits[i]], unsafe_allow_html=True)
+    text_input.render()
+    text_button = gr.Button("Predict")
 
+    with gr.Tabs():
+        with gr.TabItem("Prediction"):
+            result.render()
+           
+        with gr.TabItem("Explanation"):
+            explanation.render()
 
+    text_button.click(text_to_personality_explainer, inputs=text_input, outputs=[result, explanation])
 
-if __name__ == "__main__":
-    main()
+main.launch(show_api=False)