Update prediction.py

prediction.py

@@ -1,8 +1,14 @@
 import pandas as pd
-from tensorflow.keras.layers import TFSMLayer
+import numpy
+# from tensorflow.keras.layers import TFSMLayer
+from transformers import TFAutoModelForSequenceClassification, AutoTokenizer
 import streamlit as st
 
-model = TFSMLayer('/itsok', call_endpoint='serving_default')
+# model = TFSMLayer('/itsok', call_endpoint='serving_default')
+model_name='saved_model.pb'
+
+model = TFAutoModelForSequenceClassification.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 
 def run():
@@ -10,7 +16,7 @@ def run():
     st.write('This is a simple web app to predict sentiment of a text using deep learning. Input your feeling below to get the prediction.')
     st.write('Trust me, I have analyzed it for you!')
 
-    text = st.text_input('Text', 'I feel so sad today')
+    texts = st.text_input('Text', 'I feel so sad today')
 
     def convert_to_label(pred):
         if pred == 0:
@@ -31,7 +37,15 @@ def run():
             return 'Unknown'
 
     if st.button("Predict Your Feeling"):
-        prediction = model.predict(text)
+        # prediction = model.predict(text)
+        inputs = tokenizer(texts, return_tensors="tf", padding=True, truncation=True)
+        # Perform inference
+        outputs = model(inputs)
+        logits = outputs.logits
+        
+        # If you want to get the predicted classes
+        prediction = tf.argmax(logits, axis=-1).numpy()
+
         label = convert_to_label(prediction)
         if label == 'Normal':
             st.success('Hi! Keep up the good work! You are feeling Okay today.')