Update app.py

app.py

@@ -1,12 +1,17 @@
 import gradio as gr
 import tensorflow as tf
-from transformers import AutoTokenizer
-from transformers import TFAutoModelForSequenceClassification
-tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
+# from transformers import AutoTokenizer
+# from transformers import TFAutoModelForSequenceClassification
+# tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
+# Load model directly
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+tokenizer = AutoTokenizer.from_pretrained("ankush-003/nosqli_identifier")
+model = AutoModelForSequenceClassification.from_pretrained("ankush-003/nosqli_identifier")
 
 def predict(payload, malitious):
     inputs = tokenizer(payload, return_tensors="tf")
-    model = TFAutoModelForSequenceClassification.from_pretrained("models/ankush-003/nosqli_identifier")
+    # model = TFAutoModelForSequenceClassification.from_pretrained("ankush-003/nosqli_identifier")
     logits = model(**inputs).logits
     predicted_class_id = int(tf.math.argmax(logits, axis=-1)[0])
     # print(model.config.id2label[predicted_class_id])