Update app.py

app.py

@@ -1,29 +1,19 @@
 import streamlit as st
-import transformers
 import pandas as pd
+from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
 
-from transformers import AutoTokenizer, AutoModelForSequenceClassification, TextClassificationPipeline
-
-# Load the pre-trained BERT model
+# Load the pre-trained BERT model and pipeline for text classification
 model_name = 'nlptown/bert-base-multilingual-uncased-sentiment'
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
-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
-
+classifier = pipeline("text-classification", model=model, tokenizer=tokenizer)
 
 # Define the Streamlit app
 def app():
-    # Create a persistent DataFrame
+    # Create a persistent DataFrame to store the classification results
     if 'results' not in st.session_state:
         st.session_state.results = pd.DataFrame(columns=['text', 'toxicity', 'score'])
-    
+
     # Create a form for users to enter their text
     with st.form(key='text_form'):
         text_input = st.text_input(label='Enter your text:')
@@ -31,18 +21,25 @@ def app():
 
     # Classify the text and display the results
     if submit_button and text_input != '':
-        label, score = classify_toxicity(text_input)
+        # Classify the text using the pre-trained BERT model
+        result = classifier(text_input)[0]
+        label = result['label']
+        score = result['score']
+
+        # Display the classification results
         st.write('Classification Result:')
         st.write(f'Text: {text_input}')
         st.write(f'Toxicity: {label}')
-        st.write(f'Score: {score}')
-        
+        st.write(f'Score: {score:.3f}')
+
         # 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 __name__ == '__main__':
-    app()
+    app()