import gradio as gr
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")

def predict(payload, malitious):
    inputs = tokenizer(payload, return_tensors="tf")
    model = TFAutoModelForSequenceClassification.from_pretrained("models/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])
    expected = "Malitious" if malitious else "Benign"
    return model.config.id2label[predicted_class_id], expected 

demo = gr.Interface(
    fn=predict,
    inputs=["text","checkbox"],
    outputs=["text","text"]
)
demo.launch(debug=True)
# gr.Interface.load("models/ankush-003/nosqli_identifier").launch()