import gradio as gr
from transformers import pipeline

# Cargar el modelo de detección NSFW
nsfw_detector = pipeline("text-classification", model="ezb/NSFW-Prompt-Detector")

def detect_nsfw(text):
    # Realizar la predicción
    result = nsfw_detector(text)
    # Devolver la etiqueta y la puntuación
    return result[0]['label'], result[0]['score']

# Definir la interfaz de Gradio
iface = gr.Interface(
    fn=detect_nsfw,
    inputs=gr.inputs.Textbox(lines=2, placeholder="Introduce el texto aquí..."),
    outputs=["text", "number"],
    title="Detector de Texto NSFW",
    description="Introduce un texto para verificar si contiene contenido NSFW."
)

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