Update app.py

app.py

@@ -1,20 +1,21 @@
 import streamlit as st
 import transformers
 import torch
+import pandas as pd
 
-st.title("Sentiment Analysis App")
+st.title("Toxicity Classification App")
 
 # Set default text
 default_text = "Cody Jiang is a fantastic student in CS-UY-4613!"
 
 # Select pretrained model
-model_names = ['model', 'distilbert-base-uncased-finetuned-sst-2-english', 'bert-base-uncased', 'roberta-base']
+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
-if model_name == 'model':
-    tokenizer = transformers.DistilBertTokenizerFast.from_pretrained('./model/')
-    model = transformers.DistilBertForSequenceClassification.from_pretrained('./model/')
+if model_name == 'Codys Finetuning Language Model':
+    tokenizer = transformers.AutoTokenizer.from_pretrained('./model/')
+    model = transformers.AutoModelForSequenceClassification.from_pretrained('./model/')
 else:
     tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
     model = transformers.AutoModelForSequenceClassification.from_pretrained(model_name)
@@ -22,10 +23,10 @@ else:
 # Set device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Define sentiment labels
-sentiment_labels = ['Negative', 'Positive']
+# Define toxicity labels
+toxicity_labels = ['toxic', 'severe_toxic', 'obscene', 'threat', 'insult', 'identity_hate']
 
-# Analyze sentiment
+# Analyze toxicity
 if st.button("Submit"):
     with st.spinner("Analyzing..."):
         text = st.text_input("Enter text to analyze", default_text)
@@ -33,7 +34,8 @@ if st.button("Submit"):
         inputs = inputs.to(device)
         outputs = model(**inputs)
         logits = outputs.logits
-        predictions = torch.argmax(logits, dim=1).cpu().numpy()
-        sentiment = sentiment_labels[predictions[0]]
+        predictions = torch.sigmoid(logits).cpu().detach().numpy()[0]
+        results = pd.DataFrame({'Toxicity Class': toxicity_labels, 'Probability': predictions})
+        results = results.sort_values('Probability', ascending=False)
 
-    st.success(f"Sentiment: {sentiment}")
+    st.success(f"Toxicity Probabilities: \n{results.to_string(index=False)}")