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")
from transformers import pipeline

classifier = pipeline("sentiment-analysis", model="ankush-003/nosqli_identifier")
# classifier(payload)

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])
    prediction = classifier(payload_text)[0]    
    
    return payload_text, {prediction["label"]: prediction["score"]}

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

demo = gr.Interface(
    title="NoSQLi Detector",
    description="DistilBERT-based NoSQL Injection Payload Detection Model",
    fn=predict,
    inputs=input_elements,
    outputs=[gr.Textbox(label="Generated Payload"), gr.Label(label="Scores")]
)
demo.launch(debug=True)
# gr.Interface.load("models/ankush-003/nosqli_identifier").launch()