Update app.py

app.py

@@ -4,69 +4,65 @@ 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")
+MODEL_PATH = "bert-base-uncased"
+tokenizer = BertTokenizer.from_pretrained(MODEL_PATH)
+model = BertForSequenceClassification.from_pretrained(MODEL_PATH)
 model.eval()
 
-# Create a persistent DataFrame to store classification results
-df_results = pd.DataFrame(columns=['Text', 'Toxic', 'Obscene', 'Insult', 'Identity Hate', 'Threat', 'Severe Toxic'])
+# Create DataFrame to store classification results
+df = pd.DataFrame(columns=["Text", "Toxic", "Obscene", "Insult", "Identity Hate", "Threat", "Severe Toxic"])
 
 def classify_text(text):
     # Tokenize input text
-    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
+    encoded_input = tokenizer.encode_plus(
+        text,
+        max_length=512,
+        padding='max_length',
+        truncation=True,
+        return_tensors='pt'
+    )
+    input_ids = encoded_input['input_ids']
+    attention_mask = encoded_input['attention_mask']
+
+    # Predict toxicity using pre-trained BERT model
     with torch.no_grad():
-        outputs = model(input_ids)
+        outputs = model(input_ids, attention_mask)
         logits = outputs.logits
-        probabilities = torch.sigmoid(logits)
-    return probabilities
+        probabilities = torch.softmax(logits, dim=1)
+        toxic_prob = probabilities[0][1].item()
+        obscene_prob = probabilities[0][2].item()
+        insult_prob = probabilities[0][3].item()
+        identity_hate_prob = probabilities[0][4].item()
+        threat_prob = probabilities[0][5].item()
+        severe_toxic_prob = probabilities[0][6].item()
+
+    # Return predicted probabilities
+    return toxic_prob, obscene_prob, insult_prob, identity_hate_prob, threat_prob, severe_toxic_prob
 
 # Streamlit app
 def main():
     st.title("Toxicity Classification App")
+    text = st.text_area("Enter text to classify", height=200)
+    if st.button("Classify"):
+        if text.strip() == "":
+            st.warning("Please enter some text.")
+        else:
+            toxic_prob, obscene_prob, insult_prob, identity_hate_prob, threat_prob, severe_toxic_prob = classify_text(text)
+            st.write("Toxicity probabilities:")
+            st.write("Toxic:", toxic_prob)
+            st.write("Obscene:", obscene_prob)
+            st.write("Insult:", insult_prob)
+            st.write("Identity Hate:", identity_hate_prob)
+            st.write("Threat:", threat_prob)
+            st.write("Severe Toxic:", severe_toxic_prob)
 
-    # 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()
-
-        # 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}")
-
-        # 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.")
+            if st.button("Add to DataFrame"):
+                df.loc[len(df)] = [text, toxic_prob, obscene_prob, insult_prob, identity_hate_prob, threat_prob, severe_toxic_prob]
+                st.success("Classification results added to DataFrame.")
 
-    # Display persistent DataFrame with classification results
-    if not df_results.empty:
-        st.write("### Classification Results DataFrame")
-        st.write(df_results)
+    # Display DataFrame with classification results
+    st.subheader("Classification Results")
+    st.dataframe(df)
 
 if __name__ == "__main__":
     main()