Update app.py

app.py

@@ -1,53 +1,72 @@
 import streamlit as st
 import pandas as pd
-import transformers
 from transformers import BertTokenizer, BertForSequenceClassification
 import torch
 
 # Load pre-trained BERT model and tokenizer
-tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
-model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
+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'])
+df_results = pd.DataFrame(columns=['Text', 'Toxic', 'Obscene', 'Insult', 'Identity Hate', 'Threat', 'Severe Toxic'])
 
 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']
-
-    # Perform inference with BERT model
+    input_ids = tokenizer.encode(text, add_special_tokens=True)
+    # Convert tokenized input to tensors
+    input_ids = torch.tensor([input_ids])
+    # Feed input to the BERT model
     with torch.no_grad():
-        outputs = model(input_ids, attention_mask=attention_mask)
+        outputs = model(input_ids)
         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)
+        probabilities = torch.sigmoid(logits)
+    return probabilities
 
 # Streamlit app
 def main():
-    st.title('Toxicity Classification App')
+    st.title("Toxicity Classification App")
+
+    # Get user input
+    text = st.text_area("Enter text", "")
+    if text:
+        # Classify text using BERT model
+        probabilities = classify_text(text)
+        toxic = probabilities[0][0].item()
+        obscene = probabilities[0][1].item()
+        insult = probabilities[0][2].item()
+        identity_hate = probabilities[0][3].item()
+        threat = probabilities[0][4].item()
+        severe_toxic = probabilities[0][5].item()
 
-    # Input text box for user to enter text
-    user_text = st.text_area('Enter text:', '')
+        # Display classification results
+        st.write("### Classification Results")
+        st.write(f"Toxic: {toxic:.4f}")
+        st.write(f"Obscene: {obscene:.4f}")
+        st.write(f"Insult: {insult:.4f}")
+        st.write(f"Identity Hate: {identity_hate:.4f}")
+        st.write(f"Threat: {threat:.4f}")
+        st.write(f"Severe Toxic: {severe_toxic:.4f}")
 
-    # Button to classify text
-    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)
+        # Add classification results to DataFrame
+        if st.button("Add to Results"):
+            global df_results
+            row = {
+                'Text': text,
+                'Toxic': toxic,
+                'Obscene': obscene,
+                'Insult': insult,
+                'Identity Hate': identity_hate,
+                'Threat': threat,
+                'Severe Toxic': severe_toxic
+            }
+            df_results = df_results.append(row, ignore_index=True)
+            st.success("Classification results added to DataFrame.")
 
-    # Display classification results
-    st.header('Classification Results')
-    st.dataframe(results_df)
+    # Display persistent DataFrame with classification results
+    if not df_results.empty:
+        st.write("### Classification Results DataFrame")
+        st.write(df_results)
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()