import gradio as gr
from transformers import pipeline

# Carga un modelo ligero que no sea muy robusto
model = pipeline("text-generation", model="EleutherAI/gpt-neo-125M")

# Define una función para procesar las solicitudes
def process_prompt(prompt):
    response = model(prompt, max_length=50, num_return_sequences=1, do_sample=True)
    return response[0]["generated_text"]

# Configura la interfaz de API con Gradio
interface = gr.Interface(
    fn=process_prompt, 
    inputs="text", 
    outputs="text",
    title="Hugging Face Space API",
    description="Modelo ligero basado en GPT-Neo para probar vulnerabilidades con RedTeamer."
)

# Ejecuta la app
if __name__ == "__main__":
    interface.launch(share=True)