update app.py

app.py

@@ -14,53 +14,37 @@ import plotly.express as px
 import streamlit as st
 from streamlit_option_menu import option_menu
 
-from transformers import AutoModelForSequenceClassification
-from transformers import AutoTokenizer
-from transformers import pipeline
+st. set_page_config(layout="wide")
 
+from transformers import pipeline
 
-import torch
 import pandas as pd
 
-@st.cache(allow_output_mutation = True)
-def init_sentiment_model():
-    MODEL = 'cardiffnlp/twitter-roberta-base-sentiment'
-    tokenizer = AutoTokenizer.from_pretrained(MODEL)
-    model = AutoModelForSequenceClassification.from_pretrained(MODEL)
-    return tokenizer, model
-
 @st.cache(allow_output_mutation = True)
 def init_text_summarization_model():
     MODEL = 'facebook/bart-large-cnn'
     pipe = pipeline("summarization", model=MODEL)
     return pipe
 
-@st.cache(allow_output_mutation = True)
-def init_topic_classification():
-    MODEL = 'deepset/roberta-base-squad2'
-    pipe = pipeline('question-answering', model=MODEL, tokenizer=MODEL)
-    return pipe
-
 @st.cache(allow_output_mutation = True)
 def init_zsl_topic_classification():
-#    MODEL = 'typeform/distilbert-base-uncased-mnli'
     MODEL = 'facebook/bart-large-mnli'
     pipe = pipeline("zero-shot-classification", model=MODEL)
     template = "This text is about {}."
     return pipe, template
 
 # Model initialization    
-#tokenizer_sentiment, model_sentiment = init_sentiment_model()
 pipeline_summarization = init_text_summarization_model()
-#pipeline_topic = init_topic_classification()
-pipeline_topic, template = init_zsl_topic_classification()
+pipeline_zsl, template = init_zsl_topic_classification()
 
 st.header('Customer Review Analysis')
+
+# Review text box
 default_review = \
     """I attempted to attend the Bank Alpha Eastbank branch last Friday to open a Kids Savings Account for my 3 year old son. I was informed by the Bank Alpha staffer that I was "too close to closing time" and that "I'd have to come back another time". It was about 20 minutes prior to close and there was no one else in the branch except the staff. I did say that I had my son's birth certificate and a Medicare card to establish his identity. Still, no dice. Come back later. No worries, can do.
     I returned today (Monday) to the same branch at 1130, with more than enough time for the account opening to occur. I confirmed with another Bank Alpha staffer that I had my son's birth certificate and Medicare card. However, he went out the back "just to check something". Upon coming back, I was informed that they would not be able to open the account for me today as they required my son, a 3 year old, to be present. The staffer on Friday failed to mention this to me. Equally, the Bank Alpha website for the Kids Savings Account does not list the physical presence of the child as a requirement for opening said account.
     I have never come across a bank so committed to not providing services to prospective customers as Bank Alpha. As a result, my son won't be banking with Bank Alpha, and I probably won't be recommending the use of Bank Alpha to any family or friends either."""
-review = st.text_area("Paste/write a review here..", value=default_review, height=450)
+review = st.text_area("Paste/write a review here..", value=default_review, height=250)
 
 tabs = option_menu(menu_title=None,
                    options=[
@@ -71,36 +55,86 @@ tabs = option_menu(menu_title=None,
                    orientation='horizontal'
                    )
 
+### Text Summarization
 if tabs == 'Text Summarization':
     button = st.button('Summarize review')
     
     if button:
-        summary_text = pipeline_summarization(review, max_length=130, min_length=30, do_sample=False)
+        # Text summarization inference
+        with st.spinner("Summarizing review..."):
+            summary_text = pipeline_summarization(review, max_length=130, min_length=30, do_sample=False)
+
+        # Show output
         st.write(summary_text[0]['summary_text'])
 
+### Zero-Shot-Learning
 elif tabs == 'Zero-Shot-Learning':
-    topics = st.multiselect("",["Excellent Customer Service", "Great Product Feature", "Poor Service","Unclear Procedure"],
-                                ["Excellent Customer Service", "Great Product Feature", "Poor Service","Unclear Procedure"])
-    additional_topics = st.text_input("Enter comma separated reasons.")
-    if additional_topics:
-        topics = topics + additional_topics.split(",")
+    col_product, col_topic = st.columns(2)
     
-    is_multi_label =  st.checkbox("Can have more than one classes", value=True)
+    # Set product classes
+    products = col_product.multiselect(
+        label='Available Products and Services:', 
+        options=[
+            'Bank Account',
+            'Credit Card',
+            'Home Loan',
+            'Insurance',
+        ],
+        default=[
+            'Bank Account',
+            'Credit Card',
+            'Home Loan',
+            'Insurance',
+        ]
+    )
+    product_is_multi_label =  col_product.checkbox("Can have more than one classes", value=True)
+
+    # Set topic classes
+    topics = col_topic.multiselect(
+        label="Possible Review Topics:",
+        options=[
+            "Excellent Customer Service", 
+            "Great Product Feature", 
+            "Poor Service",
+            "Unclear Procedure",
+            "Other"
+        ],
+        default=[
+            "Excellent Customer Service", 
+            "Great Product Feature", 
+            "Poor Service",
+            "Unclear Procedure",
+        ]
+    )
+    topic_is_multi_label =  col_topic.checkbox("Can have more than one classes", value=False)
     
     button = st.button('Classify')
     
     if button:
-        classification_output = pipeline_topic(review, topics, hypothesis_template=template, multi_label=is_multi_label)
-        data = {'Class': classification_output['labels'], 'Scores': classification_output['scores']}
+        # ZSL inference
+        with st.spinner("Identifying product/service and classifying review..."):
+            product_classification_output = pipeline_zsl(review, products, hypothesis_template=template, multi_label=product_is_multi_label)
+            topic_classification_output = pipeline_zsl(review, topics, hypothesis_template=template, multi_label=topic_is_multi_label)
+
+        # Show output
+        col_output_product, col_output_topic = st.columns(2)
+
+        data = {
+            'Product': product_classification_output['labels'], 
+            'Scores': product_classification_output['scores']
+        }
         df = pd.DataFrame(data)
-        df = df.sort_values(by='Scores', ascending=False)
-        fig = px.bar(df, x='Scores', y='Class', orientation='h', width=400, height=300)
-        fig.update_layout(
-            yaxis=dict(
-                autorange='reversed'
-            )
-        )
-        st.write(fig)
-
+        df = df.sort_values(by='Scores', ascending=True)
+        fig = px.bar(df, x='Scores', y='Product', orientation='h')
+        col_output_product.plotly_chart(fig, use_container_width=True)
+
+        data = {
+            'Topic': topic_classification_output['labels'], 
+            'Scores': topic_classification_output['scores']
+        }
+        df = pd.DataFrame(data)
+        df = df.sort_values(by='Scores', ascending=True)
+        fig = px.bar(df, x='Scores', y='Topic', orientation='h')
+        col_output_topic.plotly_chart(fig, use_container_width=True)