Update app.py

app.py

@@ -1,48 +1,57 @@
 import streamlit as st
-import transformers
 import pandas as pd
+import torch
+from transformers import BertTokenizer, BertForSequenceClassification
 
-from transformers import AutoTokenizer, AutoModelForSequenceClassification, TextClassificationPipeline
+# Load the pre-trained BERT model and tokenizer
+model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)
+tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
 
-# 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)
-pipeline = TextClassificationPipeline(model=model, tokenizer=tokenizer, framework='pt', task='text-classification')
-
-# Define the toxicity classification function
+# Define the classification function
 def classify_toxicity(text):
-    result = pipeline(text)[0]
-    label = result['label']
-    score = result['score']
-    return label, score
-
+    # Tokenize the text
+    input_ids = tokenizer.encode(text, add_special_tokens=True)
+    # Truncate or pad the text to the maximum length
+    max_length = 128
+    input_ids = input_ids[:max_length] + [0] * (max_length - len(input_ids))
+    # Convert the input to a PyTorch tensor
+    input_ids = torch.tensor(input_ids).unsqueeze(0)
+    # Set the model to evaluation mode
+    model.eval()
+    # Make a prediction
+    with torch.no_grad():
+        outputs = model(input_ids)
+        logits = outputs[0]
+        # Get the predicted class
+        predicted_class = torch.argmax(logits).item()
+    return predicted_class
 
 # Define the Streamlit app
-def app():
-    # Create a persistent DataFrame
-    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:')
-        submit_button = st.form_submit_button(label='Classify')
-
-    # Classify the text and display the results
-    if submit_button and text_input != '':
-        label, score = classify_toxicity(text_input)
-        st.write('Classification Result:')
-        st.write(f'Text: {text_input}')
-        st.write(f'Toxicity: {label}')
-        st.write(f'Score: {score}')
-        
-        # 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)
-
-    # Display the persistent DataFrame
-    st.write('Classification Results:')
-    st.write(st.session_state.results)
-
-if _name_ == '_main_':
-    app()
+def main():
+    # Set the app title
+    st.title("Toxicity Classifier")
+    # Create a text input for the user to enter the text to classify
+    text = st.text_input("Enter text to classify:")
+    # Create a button to perform the classification
+    if st.button("Classify"):
+        # Classify the text
+        predicted_class = classify_toxicity(text)
+        # Display the predicted class
+        if predicted_class == 0:
+            st.write("The text is not toxic.")
+        else:
+            st.write("The text is toxic.")
+        # Create a checkbox to allow the user to add the classification to a DataFrame
+        add_to_dataframe = st.checkbox("Add classification to DataFrame")
+        if add_to_dataframe:
+            # Load the existing DataFrame or create a new one if it doesn't exist
+            dataframe = pd.read_csv("classifications.csv") if "classifications.csv" in st.session_state else pd.DataFrame(columns=["Text", "Toxic"])
+            # Add the classification to the DataFrame
+            dataframe = dataframe.append({"Text": text, "Toxic": predicted_class}, ignore_index=True)
+            # Save the DataFrame to a CSV file
+            dataframe.to_csv("classifications.csv", index=False)
+            # Display the DataFrame
+            st.write(dataframe)
+
+if __name__ == "__main__":
+    main()