Update app.py

app.py

@@ -1,53 +1,42 @@
 import streamlit as st
 import pandas as pd
-import transformers
 from transformers import BertTokenizer, BertForSequenceClassification
 import torch
 
-# Load pre-trained BERT model and tokenizer
+# Load pre-trained BERT model
 tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
 model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
 model.eval()
 
 # Create a persistent DataFrame to store classification results
-results_df = pd.DataFrame(columns=['Text', 'Toxicity'])
+classified_data = pd.DataFrame(columns=['Text', 'Toxicity'])
 
 def classify_text(text):
-    # Tokenize input text
-    inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True)
-    input_ids = inputs['input_ids']
-    attention_mask = inputs['attention_mask']
+    # Tokenize and encode input text
+    input_ids = torch.tensor(tokenizer.encode(text, add_special_tokens=True)).unsqueeze(0)
 
-    # Perform inference with BERT model
-    with torch.no_grad():
-        outputs = model(input_ids, attention_mask=attention_mask)
-        logits = outputs.logits
-        probabilities = torch.softmax(logits, dim=1)
-        toxicity_score = probabilities[0][1].item()  # Extract toxicity score
-
-    return toxicity_score
-
-def add_to_results(text, toxicity):
-    global results_df
-    results_df = results_df.append({'Text': text, 'Toxicity': toxicity}, ignore_index=True)
+    # Forward pass through BERT model
+    outputs = model(input_ids)
+    logits = outputs.logits
+    predicted_class = torch.argmax(logits, dim=1).item()
+    toxicity = "Toxic" if predicted_class == 1 else "Non-Toxic"
+    return toxicity
 
 # Streamlit app
 def main():
-    st.title('Toxicity Classification App')
-
-    # Input text box for user to enter text
-    user_text = st.text_area('Enter text:', '')
-
-    # Button to classify text
-    if st.button('Classify'):
+    st.title("Toxicity Classifier")
+    user_text = st.text_area("Enter text to classify:")
+    if st.button("Classify"):
         if user_text:
-            toxicity_score = classify_text(user_text)
-            st.write('Toxicity Score:', toxicity_score)
-            add_to_results(user_text, toxicity_score)
-
-    # Display classification results
-    st.header('Classification Results')
-    st.dataframe(results_df)
-
-if _name_ == '_main_':
-    main()
+            toxicity = classify_text(user_text)
+            st.write(f"Predicted Toxicity: {toxicity}")
+            # Add classification results to the persistent DataFrame
+            global classified_data
+            classified_data = classified_data.append({'Text': user_text, 'Toxicity': toxicity}, ignore_index=True)
+        else:
+            st.warning("Please enter some text.")
+    if st.button("View Classified Data"):
+        st.write(classified_data)
+
+if __name__ == "__main__":
+    main()