Update app.py

app.py

@@ -1,71 +1,56 @@
 import streamlit as st
 import pandas as pd
-import torch
 from transformers import BertTokenizer, BertForSequenceClassification
+import torch
 
 # 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)
+tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=6)
 model.eval()
 
-# Create DataFrame to store classification results
-df_results = pd.DataFrame(columns=['Text', 'Toxic', 'Severe Toxic', 'Obscene', 'Threat', 'Insult', 'Identity Hate'])
-
+# Function to classify text using the pre-trained BERT model
 def classify_text(text):
-    # Tokenize text
-    tokens = tokenizer.encode_plus(
-        text,
-        max_length=512,
-        truncation=True,
-        padding=True,
-        return_attention_mask=True,
-        return_tensors='pt'
-    )
-    
-    # Get model's predictions
+    # Tokenize input text
+    input_ids = tokenizer.encode(text, add_special_tokens=True)
+    # Convert tokenized input to tensor
+    input_tensor = torch.tensor([input_ids])
+    # Get model predictions
     with torch.no_grad():
-        outputs = model(**tokens)
-        logits = outputs.logits
-        probabilities = torch.softmax(logits, dim=1).tolist()[0]
-    
-    # Extract predicted labels
-    threshold = 0.5
-    labels = ['Toxic', 'Severe Toxic', 'Obscene', 'Threat', 'Insult', 'Identity Hate']
-    predicted_labels = [labels[i] for i, prob in enumerate(probabilities) if prob > threshold]
-    
+        logits = model(input_tensor)[0]
+    # Get predicted labels
+    predicted_labels = torch.sigmoid(logits).numpy()
     return predicted_labels
 
+# Create a persistent DataFrame to store classification results
+results_df = pd.DataFrame(columns=['Text', 'Toxic', 'Severe Toxic', 'Obscene', 'Threat', 'Insult', 'Identity Hate'])
+
 # Streamlit app
-def main():
-    st.title('Toxicity Classification')
+def app():
+    st.title("Toxicity Classification App")
+    st.write("Enter text below to classify its toxicity.")
     
     # User input
-    text = st.text_area('Enter text:', max_chars=512)
+    user_input = st.text_area("Enter text here:", "", key='user_input')
     
-    # 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.')
+    # Classification
+    if st.button("Classify"):
+        # Perform classification
+        labels = classify_text(user_input)
+        # Print classification results
+        st.write("Classification Results:")
+        st.write("Toxic: {:.2%}".format(labels[0][0]))
+        st.write("Severe Toxic: {:.2%}".format(labels[0][1]))
+        st.write("Obscene: {:.2%}".format(labels[0][2]))
+        st.write("Threat: {:.2%}".format(labels[0][3]))
+        st.write("Insult: {:.2%}".format(labels[0][4]))
+        st.write("Identity Hate: {:.2%}".format(labels[0][5]))
+        # Add results to persistent DataFrame
+        results_df.loc[len(results_df)] = [user_input, labels[0][0], labels[0][1], labels[0][2], labels[0][3], labels[0][4], labels[0][5]]
     
-    # Show DataFrame with classification results
-    if not df_results.empty:
-        st.subheader('Classification Results')
-        st.dataframe(df_results)
+    # Show results DataFrame
+    st.write("Classification Results DataFrame:")
+    st.write(results_df)
 
-if __name__ == '__main__':
-    main()
+# Run the app
+if __name__ == "__main__":
+    app()