# Importar las bibliotecas necesarias
from transformers import pipeline
import gradio as gr

# Crear la pipeline de generación de texto con el modelo español especificado
pipe = pipeline("text-generation", model="datificate/gpt2-small-spanish")

# Función que el chatbot utilizará para responder a los prompts
def chatbot_response(prompt):
    # Genera la respuesta del modelo, activando explícitamente el truncamiento
    result = pipe(prompt, max_length=50, truncation=True)
    # Retorna solo el texto de la respuesta generada
    return result[0]['generated_text']

# Crear la interfaz de Gradio
iface = gr.Interface(
    fn=chatbot_response, 
    inputs="text", 
    outputs="text",
    title="Chatbot en Español",
    description="Este chatbot responde a tus preguntas. Está basado en el modelo GPT-2 pequeño en español.",
    theme="default"  # Puedes cambiar el tema de la interfaz si lo deseas
)

# Lanzar la interfaz, activando la opción para compartir si se desea un enlace público
iface.launch(share=True)