Update app.py

app.py

@@ -1,68 +1,58 @@
 import streamlit as st
 import pandas as pd
 from transformers import BertTokenizer, BertForSequenceClassification
-import torch
 
 # Load pre-trained BERT model and tokenizer
-MODEL_PATH = "bert-base-uncased"
-tokenizer = BertTokenizer.from_pretrained(MODEL_PATH)
-model = BertForSequenceClassification.from_pretrained(MODEL_PATH)
+tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=6)
+model.load_state_dict(torch.load('path_to_pretrained_model.pth', map_location=torch.device('cpu'))) # Load pre-trained model weights
 model.eval()
 
-# Create DataFrame to store classification results
-df = pd.DataFrame(columns=["Text", "Toxic", "Obscene", "Insult", "Identity Hate", "Threat", "Severe Toxic"])
-
+# Define function for classification
 def classify_text(text):
-    # Tokenize input text
-    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, attention_mask)
-        logits = outputs.logits
-        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
+    inputs = tokenizer(text, padding=True, truncation=True, return_tensors="pt")
+    outputs = model(**inputs)
+    logits = outputs.logits
+    probabilities = torch.sigmoid(logits)
+    return probabilities
 
-# Streamlit app
+# Create Streamlit app
 def main():
     st.title("Toxicity Classification App")
-    text = st.text_area("Enter text to classify", height=200)
+    text_input = st.text_area("Enter text for classification:", "")
     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)
-
-            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 DataFrame with classification results
-    st.subheader("Classification Results")
-    st.dataframe(df)
+        probabilities = classify_text(text_input)
+        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()
+
+        st.write("### Classification Results:")
+        st.write("Toxic: {:.2%}".format(toxic))
+        st.write("Obscene: {:.2%}".format(obscene))
+        st.write("Insult: {:.2%}".format(insult))
+        st.write("Identity Hate: {:.2%}".format(identity_hate))
+        st.write("Threat: {:.2%}".format(threat))
+        st.write("Severe Toxic: {:.2%}".format(severe_toxic))
+
+        # Add results to persistent DataFrame
+        result = {
+            'Text': text_input,
+            'Toxic': toxic,
+            'Obscene': obscene,
+            'Insult': insult,
+            'Identity Hate': identity_hate,
+            'Threat': threat,
+            'Severe Toxic': severe_toxic
+        }
+        df = pd.DataFrame(result, index=[0])
+        if 'classified_df' not in st.session_state:
+            st.session_state.classified_df = pd.DataFrame(columns=['Text', 'Toxic', 'Obscene', 'Insult', 'Identity Hate', 'Threat', 'Severe Toxic'])
+        st.session_state.classified_df = st.session_state.classified_df.append(df, ignore_index=True)
+        st.write("### Classified Data:")
+        st.write(st.session_state.classified_df)
 
 if __name__ == "__main__":
     main()