app.py updated

README.md

@@ -4,6 +4,7 @@ emoji: 🌍
 colorFrom: blue
 colorTo: yellow
 sdk: streamlit
+sdk_version: 1.17.0
 app_file: app.py
 pinned: false
 license: apache-2.0

app.py

@@ -6,9 +6,14 @@ single = SinglePrediction()
 batch = BatchPrediction()
 
 def single_predict(text):
-    st.success(f'{text} :thumbsup:')
-    preds = single.predict(text)
-    #st.plotly_chart(preds, theme=None, use_container_width=True)
+    preds, fig = single.predict(text)
+
+    if preds < 0.5:
+        st.success(f'Non Toxic Comment!!! :thumbsup:')
+        st.plotly_chart(fig, theme="streamlit", use_container_width=True)
+    else:
+        st.error(f'Toxic Comment!!! :thumbsup:')
+        st.plotly_chart(fig, theme="streamlit", use_container_width=True)
 
 def batch_predict(data):
     if batch.data_validation(data):

src/constants.py

@@ -3,8 +3,8 @@ import os
 ROOT_DIR = os.getcwd()
 MODEL_DIR_NAME = "serving_model"
 MODEL_NAME = "roberta-fine-tuned-2"
-MODEL_PATH = os.path.join(ROOT_DIR, MODEL_DIR_NAME,MODEL_NAME)
+MODEL_PATH = os.path.join(ROOT_DIR, MODEL_NAME)
 TOKENIZER_FILE_NAME = "tokenizer"
-TOKENIZER_PATH = os.path.join(ROOT_DIR, MODEL_DIR_NAME, TOKENIZER_FILE_NAME)
+TOKENIZER_PATH = os.path.join(ROOT_DIR, TOKENIZER_FILE_NAME)
 MAX_LEN =192
 BUFFER_SIZE=2048

src/single_predict.py

@@ -1,6 +1,7 @@
 import pandas as pd
 import numpy as np
 import tensorflow as tf
+import plotly.express as px
 import transformers
 from transformers import AutoTokenizer
 import os
@@ -22,12 +23,24 @@ class SinglePrediction:
                                 return_tensors="tf",
                                 return_token_type_ids = False)
         return dict(tokens)
+    
+    def plot(self, pred):
+        fig = px.bar(x=[round(pred), round(1-pred)], 
+                     y=['toxic', 'non-toxic'],
+                     width=500, height=250, 
+                     template="plotly_dark", 
+                     text_auto='1', 
+                     title="Probabilities(%)")
+        fig.update_traces(width=0.3,textfont_size=15, textangle=0, textposition="outside")
+        fig.update_layout(yaxis_title=None,xaxis_title=None)
+        return fig
 
     def predict(self, text:str):
         try:
             text = re.sub('\n',' ',text).strip()
             input = self.tokenizer(text)
-            preds = self.model.predict(input)[0][0]
-            return preds
+            pred = self.model.predict(input)[0][0]
+            fig = self.plot(pred)
+            return pred, fig
         except Exception as e:
             print(e)