Create app.py

app.py

@@ -0,0 +1,69 @@
+import streamlit as st
+import pandas as pd
+import numpy as np
+from unidecode import unidecode
+import tensorflow as tf 
+import cloudpickle
+from transformers import DistilBertTokenizerFast
+import os
+from matplotlib import pyplot as plt
+from PIL import Image
+
+
+with open(os.path.join("models", "toxic_comment_preprocessor_classnames.bin"), "rb") as model_file_obj:
+        text_preprocessor, class_names = cloudpickle.load(model_file_obj)    
+interpreter = tf.lite.Interpreter(model_path=os.path.join("models", "toxic_comment_classifier_hf_distilbert.tflite"))
+
+
+def inference(text):
+    text = text_preprocessor.preprocess(pd.Series(text))[0]
+    model_checkpoint = "distilbert-base-uncased"
+    tokenizer = DistilBertTokenizerFast.from_pretrained(model_checkpoint)
+    tokens = tokenizer(text, max_length=512, padding="max_length", truncation=True, return_tensors="tf")
+    
+    # tflite model inference  
+    interpreter.allocate_tensors()
+    input_details = interpreter.get_input_details()
+    output_details = interpreter.get_output_details()[0]
+    attention_mask, input_ids = tokens['attention_mask'], tokens['input_ids']
+    interpreter.set_tensor(input_details[0]["index"], attention_mask)
+    interpreter.set_tensor(input_details[1]["index"], input_ids)
+    interpreter.invoke()
+    tflite_pred = interpreter.get_tensor(output_details["index"])[0]
+    result_df = pd.DataFrame({'class': class_names, 'prob': tflite_pred})
+    result_df.sort_values(by='prob', ascending=True, inplace=True)
+    return result_df
+
+
+def display_image(df):
+    fig, ax = plt.subplots(figsize=(2, 1.8))
+    df.plot(x='class', y='prob', kind='barh', ax=ax, color='black', ylabel='')
+    ax.tick_params(axis='both', which='major', labelsize=8.5)
+    ax.get_legend().remove()
+    ax.spines['top'].set_visible(False)
+    ax.spines['right'].set_visible(False)
+    ax.spines['bottom'].set_visible(False)
+    ax.spines['left'].set_visible(False)
+    ax.get_xaxis().set_ticks([])
+    plt.rcParams["figure.autolayout"] = True
+    plt.xlim(0, 1)
+    for n, i in enumerate([*df['prob']]):
+        plt.text(i+0.015, n-0.15, f'{str(np.round(i, 3))}   ', fontsize=7.5)
+
+    fig.savefig("prediction.png", bbox_inches='tight', dpi=100)
+    image = Image.open('prediction.png')
+    st.write('')
+    st.image(image, output_format="PNG", caption="Prediction")
+    
+############## ENTRY POINT START #######################
+def main():
+    st.title("Toxic Comment Classifier")
+    comment_txt = st.text_area("Enter a comment:", "", height=100)
+    if st.button("Submit"):
+        df = inference(comment_txt)
+        display_image(df)
+
+############## ENTRY POINT END #######################
+
+if __name__ == "__main__":
+    main()