app.py

import gradio as gr
from transformers import pipeline

# Carga un modelo un poco más potente pero rápido
model = pipeline("text-generation", model="facebook/opt-350m")

# Define una función para procesar las solicitudes
def process_prompt(prompt):
    response = model(prompt, max_length=100, 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 - OPT-350M",
    description="Modelo ligero y rápido basado en OPT-350M para probar vulnerabilidades con RedTeamer.",
)

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