Update app.py

app.py

@@ -3,18 +3,18 @@ import transformers
 import torch
 import pandas as pd
 
-st.title("Toxicity Analysis App")
+st.title("Toxicity Classification App")
 
 # Set default text
 default_text = "Cody Jiang is a fantastic student in CS-UY-4613!"
 
-# Select model
-model_names = ['Codys Finetuning Language Model', 'distilbert-base-uncased-finetuned-sst-2-english', 'bert-base-uncased', 'roberta-base']
-model_name = st.selectbox("Select a model", model_names)
+# Select pretrained model
+model_names = ['distilbert-base-uncased-finetuned-sst-2-english', 'bert-base-uncased', 'roberta-base', 'Codys Finetuning Language Model']
+model_name = st.selectbox("Select a pretrained model", model_names)
 
-# Load selected model
+# Load selected model and tokenizer
 if model_name == 'Codys Finetuning Language Model':
-    model_path = 'model'
+    model_path = './model'
 else:
     model_path = model_name
 tokenizer = transformers.AutoTokenizer.from_pretrained(model_path)
@@ -34,11 +34,7 @@ if st.button("Submit"):
         inputs = inputs.to(device)
         outputs = model(**inputs)
         logits = outputs.logits
-        probabilities = torch.sigmoid(logits).detach().cpu().numpy().flatten()
-        predictions = [toxicity_labels[i] for i, p in enumerate(probabilities) if p > 0.5]
-        probabilities = [p for p in probabilities if p > 0.5]
-        data = {'Text': [text] * len(predictions), 'Toxicity': predictions, 'Probability': probabilities}
-        df = pd.DataFrame(data)
-
-    st.success("Toxicity Analysis Results:")
-    st.table(df)
+        probabilities = torch.sigmoid(logits).detach().cpu().numpy()[0]
+        predictions = (probabilities >= 0.5).astype(int)
+        result_df = pd.DataFrame({'Text': [text], 'Probability': probabilities, 'Label': toxicity_labels, 'Prediction': predictions})
+        st.table(result_df[['Text', 'Prediction', 'Probability']])