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.Textbox(lines=2, placeholder="Introduce the text here..."),
    outputs=[gr.Text(label="Label"), gr.Number(label="Confidence")],
    title="NSFW Text Detector",
    description="This model detects NSFW content in text."
)

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