app.py

import gradio as gr
import json
import tensorflow as tf
# from transformers import AutoTokenizer
# from transformers import TFAutoModelForSequenceClassification

# Load model directly
from transformers import AutoTokenizer, TFAutoModelForSequenceClassification

# tokenizer = AutoTokenizer.from_pretrained("ankush-003/nosqli_identifier")
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
model = TFAutoModelForSequenceClassification.from_pretrained("ankush-003/nosqli_identifier")

def predict(username, pwd, label, payload_text = None):
    if(payload_text is None or payload_text is ""):
        payload = {
            "username": username,
            "password": pwd
        }
        payload_text = json.dumps(payload)
    inputs = tokenizer(payload_text, return_tensors="tf")
    # 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])
    
    return payload_text, model.config.id2label[predicted_class_id]

input_elements = [gr.Textbox(label="Enter Username"), gr.Textbox(label="Enter Password"), gr.Dropdown(["Malitious", "Benign"], label="Expected", info="Enter expected value"),
            gr.Textbox(label="Enter Payload", info="Optional if username and password entered already")]

demo = gr.Interface(
    fn=predict,
    inputs=input_elements,
    outputs=[gr.Textbox(label="Generated Payload"), gr.Textbox(label="Model Prediction")]
)
demo.launch(debug=True)
# gr.Interface.load("models/ankush-003/nosqli_identifier").launch()