Update app.py

app.py

@@ -1,63 +1,71 @@
 import streamlit as st
 import pandas as pd
-from transformers import BertTokenizer, BertForSequenceClassification
 import torch
+from transformers import BertTokenizer, BertForSequenceClassification
 
-# Load pre-trained BERT model
-tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
-model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
+# Load pre-trained BERT model and tokenizer
+MODEL_NAME = 'bert-base-uncased'
+tokenizer = BertTokenizer.from_pretrained(MODEL_NAME)
+model = BertForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=6)
 model.eval()
 
-# Create a DataFrame to store classification results
-classification_results_df = pd.DataFrame(columns=['Text', 'Toxic', 'Severe Toxic', 'Obscene', 'Threat', 'Insult', 'Identity Hate'])
+# Create DataFrame to store classification results
+df_results = pd.DataFrame(columns=['Text', 'Toxic', 'Severe Toxic', 'Obscene', 'Threat', 'Insult', 'Identity Hate'])
 
 def classify_text(text):
     # Tokenize text
-    inputs = tokenizer(text, return_tensors='pt', truncation=True, padding=True)
-    # Forward pass through the BERT model
-    outputs = model(**inputs)
-    # Get predicted probabilities for each class
-    probs = torch.sigmoid(outputs.logits)
-    # Round probabilities to 0 or 1 to get binary predictions
-    preds = (probs > 0.5).int().tolist()[0]
-    return preds
-
-def add_classification_to_df(text, preds):
-    # Add classification results to the DataFrame
-    classification_results_df.loc[len(classification_results_df)] = [text] + preds
+    tokens = tokenizer.encode_plus(
+        text,
+        max_length=512,
+        truncation=True,
+        padding=True,
+        return_attention_mask=True,
+        return_tensors='pt'
+    )
+    
+    # Get model's predictions
+    with torch.no_grad():
+        outputs = model(**tokens)
+        logits = outputs.logits
+        probabilities = torch.softmax(logits, dim=1).tolist()[0]
+    
+    # Extract predicted labels
+    labels = ['Toxic', 'Severe Toxic', 'Obscene', 'Threat', 'Insult', 'Identity Hate']
+    predicted_labels = [labels[i] for i, prob in enumerate(probabilities) if prob > 0.5]
+    
+    return predicted_labels
 
 # Streamlit app
 def main():
-    st.title("Toxicity Classification with BERT")
-    # Input text from user
-    text = st.text_area("Enter text for classification", "")
-    if st.button("Classify"):
-        if text.strip() == "":
-            st.warning("Please enter some text for classification.")
-        else:
-            # Perform classification
-            preds = classify_text(text)
-            # Display classification results
-            st.subheader("Classification Results:")
-            # Check if preds has enough elements
-            if len(preds) >= 6:
-                st.write("Toxic: ", preds[0])
-                st.write("Severe Toxic: ", preds[1])
-                st.write("Obscene: ", preds[2])
-                st.write("Threat: ", preds[3])
-                st.write("Insult: ", preds[4])
-                st.write("Identity Hate: ", preds[5])
-                # Add classification results to DataFrame
-                add_classification_to_df(text, preds)
-            else:
-                st.error("Error: Classification results are incomplete.")
-                # Debug statements
-                st.write("preds:", preds)
-                st.write("len(preds):", len(preds))
-    if st.button("View Classification Results"):
-        # Display classification results DataFrame
-        st.subheader("All Classification Results:")
-        st.write(classification_results_df)
+    st.title('Toxicity Classification')
+    
+    # User input
+    text = st.text_area('Enter text:', max_chars=512)
+    
+    # Perform classification
+    if st.button('Classify'):
+        predicted_labels = classify_text(text)
+        st.write('Predicted Labels:', predicted_labels)
+        
+        # Allow user to add classification results to DataFrame
+        if st.button('Add to Results'):
+            global df_results
+            df_results = df_results.append({
+                'Text': text,
+                'Toxic': 'Toxic' in predicted_labels,
+                'Severe Toxic': 'Severe Toxic' in predicted_labels,
+                'Obscene': 'Obscene' in predicted_labels,
+                'Threat': 'Threat' in predicted_labels,
+                'Insult': 'Insult' in predicted_labels,
+                'Identity Hate': 'Identity Hate' in predicted_labels
+            }, ignore_index=True)
+            st.success('Classification results added to DataFrame.')
+    
+    # Show DataFrame with classification results
+    if not df_results.empty:
+        st.subheader('Classification Results')
+        st.dataframe(df_results)
+    
 
 if __name__ == '__main__':
     main()