Update app.py

app.py

@@ -1,94 +1,56 @@
 import streamlit as st
 import pandas as pd
-import numpy as np
-from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification, BertForSequenceClassification, DistilBertModel
+from transformers import BertTokenizer, BertForSequenceClassification
 import torch
-from torch import cuda
-from torch.utils.data import Dataset, DataLoader
-import finetuning
-from finetuning import CustomDistilBertClass
 
-# 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
-model_options = list(model_map.keys())
-
-# 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']
-
-@st.cache_resource
-def load_model(model_name):
-    """Load pretrained BERT model."""
-    path = "finetuned_model.pt"
-    model = torch.load(path)
-    tokenizer = AutoTokenizer.from_pretrained(model_map[model_name])
-    return model, tokenizer
-
-def classify_text(model, tokenizer, text):
-    """Classify text using pretrained BERT model."""
-    inputs = tokenizer.encode_plus(
-        text,
-        add_special_tokens=True,
-        max_length=512,
-        padding='max_length',
-        return_tensors='pt',
-        truncation=True
-    )
-    print(inputs)
-    with torch.no_grad():
-        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"pred class: {pred_class}")
-        print(probabilities[0].tolist())
-    return label_cols[pred_class], round(probabilities[0].tolist(),2)
-
-# 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
-st.subheader('Enter comment below:')
-text_input = st.text_area(label='', height=100, max_chars=500)
-
-# Make prediction when user clicks 'Classify' button
-if st.button('Classify Toxicity'):
-    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('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=['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('Tweet:', text_input)
-        st.write('Highest Toxicity Class:', class_label)
-        st.write('Probability:', class_prob)
-        st.session_state.classification_results = st.session_state.classification_results.append({
-            'tweet': text_input,
-            'toxicity_class': class_label,
-            'probability': class_prob
-        }, ignore_index=True)
-st.write(st.session_state.classification_results)
+# Load pre-trained BERT model
+tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
+model.eval()
+
+# Create a DataFrame to store classification results
+classification_results_df = pd.DataFrame(columns=['Text', 'Toxic', 'Severe Toxic', 'Obscene', 'Threat', 'Insult', 'Identity Hate'])
+
+def classify_text(text):
+    # Tokenize text
+    inputs = tokenizer(text, return_tensors='pt', truncation=True, padding=True)
+    # Forward pass through the BERT model
+    outputs = model(**inputs)
+    # Get predicted probabilities for each class
+    probs = torch.sigmoid(outputs.logits)
+    # Round probabilities to 0 or 1 to get binary predictions
+    preds = (probs > 0.5).int().tolist()[0]
+    return preds
+
+def add_classification_to_df(text, preds):
+    # Add classification results to the DataFrame
+    classification_results_df.loc[len(classification_results_df)] = [text] + preds
+
+# Streamlit app
+def main():
+    st.title("Toxicity Classification with BERT")
+    # Input text from user
+    text = st.text_area("Enter text for classification", "")
+    if st.button("Classify"):
+        if text.strip() == "":
+            st.warning("Please enter some text for classification.")
+        else:
+            # Perform classification
+            preds = classify_text(text)
+            # Display classification results
+            st.subheader("Classification Results:")
+            st.write("Toxic: ", preds[0])
+            st.write("Severe Toxic: ", preds[1])
+            st.write("Obscene: ", preds[2])
+            st.write("Threat: ", preds[3])
+            st.write("Insult: ", preds[4])
+            st.write("Identity Hate: ", preds[5])
+            # Add classification results to DataFrame
+            add_classification_to_df(text, preds)
+    if st.button("View Classification Results"):
+        # Display classification results DataFrame
+        st.subheader("All Classification Results:")
+        st.write(classification_results_df)
+
+if __name__ == '__main__':
+    main()