Update app.py

app.py

@@ -1,45 +1,62 @@
 import streamlit as st
+import transformers
 import pandas as pd
-from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
 
-# Load the pre-trained BERT model and pipeline for text classification
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, TextClassificationPipeline
+
+# Load the pre-trained BERT model
 model_name = 'nlptown/bert-base-multilingual-uncased-sentiment'
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
-classifier = pipeline("text-classification", model=model, tokenizer=tokenizer)
+pipeline = TextClassificationPipeline(model=model, tokenizer=tokenizer, framework='pt', task='text-classification')
+
+# Define the toxicity classification function
+def classify_toxicity(text):
+    result = pipeline(text)[0]
+    label = result['label']
+    score = result['score']
+    return label, score
 
 # Define the Streamlit app
 def app():
-    # Create a persistent DataFrame to store the classification results
+    # Create a persistent DataFrame
     if 'results' not in st.session_state:
         st.session_state.results = pd.DataFrame(columns=['text', 'toxicity', 'score'])
 
+    # Set page title and favicon
+    st.set_page_config(page_title='Toxicity Classifier', page_icon=':microscope:')
+
+    # Set app title and description
+    st.title('Toxicity Classifier')
+    st.markdown('This app uses a pre-trained BERT model to classify text into three toxicity levels: **positive**, **neutral**, and **negative**.')
+
     # Create a form for users to enter their text
     with st.form(key='text_form'):
-        text_input = st.text_input(label='Enter your text:')
+        text_input = st.text_input(label='Enter your text:', help='Enter some text to classify.')
         submit_button = st.form_submit_button(label='Classify')
 
     # Classify the text and display the results
     if submit_button and text_input != '':
-        # Classify the text using the pre-trained BERT model
-        result = classifier(text_input)[0]
-        label = result['label']
-        score = result['score']
+        label, score = classify_toxicity(text_input)
 
-        # Display the classification results
-        st.write('Classification Result:')
-        st.write(f'Text: {text_input}')
-        st.write(f'Toxicity: {label}')
-        st.write(f'Score: {score:.3f}')
+        # Create a collapsible section for the classification result
+        with st.beta_expander('Classification Result', expanded=True):
+            st.write(f'Text: {text_input}')
+            st.write(f'Toxicity: {label}')
+            st.write(f'Score: {score:.4f}')
 
         # Add the classification result to the persistent DataFrame
-        st.session_state.results = st.session_state.results.append(
-            {'text': text_input, 'toxicity': label, 'score': score}, ignore_index=True
-        )
+        st.session_state.results = st.session_state.results.append({'text': text_input, 'toxicity': label, 'score': score}, ignore_index=True)
 
     # Display the persistent DataFrame
-    st.write('Classification Results:')
-    st.write(st.session_state.results)
+    if not st.session_state.results.empty:
+        st.markdown('## Classification Results')
+        st.dataframe(st.session_state.results.style.highlight_max(subset=['score'], color='lightgreen'))
+
+    # Add a footer with some info and links
+    st.markdown('---')
+    st.markdown('Made with :heart: by ChatGPT')
+    st.markdown('[GitHub](https://github.com/chatgpt/toxicity-classifier) - [OpenAI Blog](https://openai.com/blog/) - [Streamlit Docs](https://docs.streamlit.io/)')
 
 if __name__ == '__main__':
     app()