import gradio as gr
from transformers import BertTokenizer, BertForSequenceClassification
from transformers import pipeline

urlbert = BertForSequenceClassification.from_pretrained('elftsdmr/malware-url-detect',num_labels=2)
tokenizer = BertTokenizer.from_pretrained('elftsdmr/malware-url-detect')

url_classifier = pipeline("text-classification", model=urlbert, tokenizer=tokenizer)





def text_to_url(text):
    return url_classifier(text)[0]["label"]


demo = gr.Blocks(theme=gr.themes.Default(primary_hue="red", secondary_hue="pink"))

with demo:
    
    text = gr.Textbox()
    label = gr.Label()

  
    b2 = gr.Button("Classify URL")
 
   
    b2.click(text_to_url, inputs=text, outputs=label)

if __name__ == "__main__":
    demo.launch()