Update app.py

app.py

@@ -1,4 +1,3 @@
-# Import necessary libraries
 import streamlit as st
 import pandas as pd
 import numpy as np
@@ -9,22 +8,22 @@ from torch.utils.data import Dataset, DataLoader
 import finetuning
 from finetuning import CustomDistilBertClass
 
-# Define the pretrained models to be used
+# device = 'cuda' if cuda.is_available() else 'cpu'
+# Load pretrained models
 model_map = {
     'BERT': 'bert-base-uncased',
     'RoBERTa': 'roberta-base',
     'DistilBERT': 'distilbert-base-uncased'
 }
 
-# Load dropdown options for model selection
+# Load dropdown options
 model_options = list(model_map.keys())
 
-# Load the training dataset
+# Load dataset
 train_df = pd.read_csv('train.csv')
 train_df = train_df.sample(n=256)
 label_cols = ['toxic', 'severe_toxic', 'obscene', 'threat', 'insult', 'identity_hate']
 
-# Define a function to load the pretrained BERT model
 @st.cache_resource
 def load_model(model_name):
     """Load pretrained BERT model."""
@@ -33,7 +32,6 @@ def load_model(model_name):
     tokenizer = AutoTokenizer.from_pretrained(model_map[model_name])
     return model, tokenizer
 
-# Define a function to classify text using the pretrained BERT model
 def classify_text(model, tokenizer, text):
     """Classify text using pretrained BERT model."""
     inputs = tokenizer.encode_plus(
@@ -49,18 +47,18 @@ def classify_text(model, tokenizer, text):
         logits = model(inputs['input_ids'],inputs['attention_mask'])[0]
         probabilities = torch.softmax(logits, dim=1)[0]
         pred_class = torch.argmax(probabilities, dim=0)
-        print(f"predicted class: {pred_class}")
+        print(f"pred class: {pred_class}")
         print(probabilities[0].tolist())
     return label_cols[pred_class], round(probabilities[0].tolist(),2)
 
-# Set up Streamlit app
+# Set up streamlit app
 st.title('Toxic Comment Classifier')
 model_name = st.sidebar.selectbox('Select a model', model_options)
 st.sidebar.write('Selected:', model_name)
 model, tokenizer = load_model(model_name)
 print(type(model))
 
-# Define input text area for user input
+# Define input text area
 st.subheader('Enter comment below:')
 text_input = st.text_area(label='', height=100, max_chars=500)
 
@@ -71,21 +69,25 @@ if st.button('Classify Toxicity'):
     else:
         class_label, class_prob = classify_text(model, tokenizer, text_input)       
         st.subheader('Results')
-        st.write('Comment:', text_input)
+        st.write('Tweet:', text_input)
         st.write('Highest Toxicity Class:', class_label)
         st.write('Probability:', class_prob)
 
 # Display table of results
 st.subheader('Toxic Classification Results')
 if 'classification_results' not in st.session_state:
-    st.session_state.classification_results = pd.DataFrame(columns=['comment', 'toxicity_class', 'probability'])
+    st.session_state.classification_results = pd.DataFrame(columns=['tweet', 'toxicity_class', 'probability'])
 if st.button('Add to Results'):
     if not text_input:
         st.write('Please enter comment')
     else:
         class_label, class_prob = classify_text(model, tokenizer, text_input)       
         st.subheader('Results')
-        st.write('Comment:', text_input)
+        st.write('Tweet:', text_input)
         st.write('Highest Toxicity Class:', class_label)
         st.write('Probability:', class_prob)
-        st.session_state.classification_results = st.session_state.classification
+        st.session_state.classification_results = st.session_state.classification_results.append({
+            'tweet': text_input,
+            'toxicity_class': class_label,
+            'probability': class_prob
+        },