app.py

import gradio as gr
from transformers import pipeline
from transformers import AutoTokenizer


## Cargando el modelo ##
def load_model():
    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")  # Use the base model's tokenizer
    return pipeline("text-classification", model="Fredin14/bert-base-cased-finetuned-PhishingClassificationl", tokenizer=tokenizer,  padding=True, truncation=True)
model = load_model()

## Funcion para hacer las predicciones ##
def predict(text_input):
    predictions = model(text_input)
    # Convertir las etiquetas a 1 (Phishing) o 0 (Benign)
    if "LABEL_1" == predictions[0]['label']:
        label = "Phishing"
    else:
        label = "Bening"

    confidence = predictions[0]['score']
    result_str = f"{label} (Confidence: {confidence:.2%})"
    
    return text_input, result_str


## Interfaz para el usuario ##
phishing_app = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(
        label="Enter text to analyze", 
        placeholder="E.g.: 'Click here to claim your prize...'",
        lines=5
    ),
    outputs=[
        gr.Textbox(label="Analyzed text"), 
        gr.Label(label="Result", num_top_classes=2)
    ],
    title="Phishing Detection Tool",
    description="Analyze whether a text is potentially phishing",
    examples=[
        ["Urgent: Your account will be suspended unless you verify your details now"],
        ["Dear customer, your invoice #45612 is ready for payment"],
        ["You won an iPhone 15! Click to claim your prize"],
        ["Security alert: Unusual login attempt detected on your account"],
        ["hey, I wondering if you are free tonight"]
    ],
    allow_flagging="manual",
    flagging_options=["Flag as false positive", "Flag as interesting"]
)


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