Update app.py

app.py

@@ -1,45 +1,24 @@
 import streamlit as st
 from transformers import pipeline
+import tensorflow as tf
+import numpy as np
+from tensorflow.keras.layers import TextVectorization
+from tensorflow import keras
 
-st.title('Sentiment Analyser')
-st.image("https://www.pngall.com/wp-content/uploads/5/Emotion-Transparent.png")
+model = tf.keras.models.load_model('toxicity_model.h5')
 
-st.write("---")
-st.header('Step 1 - Select a language model')
-col1, col2 = st.columns(2)
+df = pd.read_csv('train.csv')
+X = df['comment_text']
+y = df[df.columns[2:]].values
 
-selection = "N/A"
-with col1:
-    selection = st.radio("Pick one of the four pre-trained models below:",
-        key = "modelChoice",
-        options = ["DistilBERT", "Toxicity-Classifier", "SiEBERT", "Twitter-roBERTa"],
-    )
+MAX_FEATURES = 200000
+vectorizer = TextVectorization(max_tokens=MAX_FEATURES,
+                               output_sequence_length=1800,
+                               output_mode='int')
 
-pipe = pipeline('sentiment-analysis')
-if selection == "DistilBERT":
-    pipe = pipeline(model = "distilbert-base-uncased-finetuned-sst-2-english")
-if selection == "Toxicity-Classifier":
-    pipe = pipeline(model = "JulianHame/Toxicity-Classifier")
-if selection == "Twitter-roBERTa":
-    pipe = pipeline(model = "unitary/toxic-bert")
-if selection == "SiEBERT":
-    pipe = pipeline(model = "siebert/sentiment-roberta-large-english")
+vectorizer.adapt(X.values)
 
-with col2:
-    st.caption('DistilBERT - One of the most popular and widely-used language models. Labels text as POSITIVE or NEGATIVE. Developed by Hugging Face.')
-    st.caption('Toxicity-Classifier - A model trained to classify tweets under different toxicity-related categories.')
-    st.caption('SiEBERT - A model trained on diverse text sources to improve generalization. Labels text as POSITIVE or NEGATIVE. Developed by siebert.')
-    st.caption('Twitter-roBERTa - A model trained on over 124M tweets. Labels text as POSITIVE, NEGATIVE or NEUTRAL. Developed by cardiffnlp.')
-
-st.write("---")
-st.header('Step 2 - Enter some text')
-text = st.text_area('The sentiment of the text entered here will be determined based on the model you chose above.',
-                   value = "It was the best of times, it was the worst of times.")
-
-st.write("---")
-st.header('Step 3 - View your results')
-
-if text:
-    st.write('Model used: ', selection)
-    out = pipe(text)
-    st.json(out)
+input_str = vectorizer('I hate you.')
+res = model.predict(np.expand_dims(input_str,0))
+classification = res[0].tolist()
+st.write(classification)