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, payload_text = None, label):
    if(payload_text is None):
        payload["username"] = username
        payload["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 model.config.id2label[predicted_class_id]

demo = gr.Interface(
    fn=predict,
    inputs=[gr.Textbox(label="Enter Username"),gr.Textbox(label="Enter Password"),
            gr.Textbox(label="Enter Payload[optional]"),
           gr.Dropdown(
            ["Malitious", "Benign"], label="Expected", info="Enter expected value"
        )],
    outputs=[gr.Textbox(label="Model Prediction")]
)
demo.launch(debug=True)
# gr.Interface.load("models/ankush-003/nosqli_identifier").launch()