import gradio as gr
from transformers import GPT2LMHeadModel, GPT2Tokenizer

MODEL_NAME = "datificate/gpt2-small-spanish"

tokenizer = GPT2Tokenizer.from_pretrained(MODEL_NAME)
model = GPT2LMHeadModel.from_pretrained(MODEL_NAME)

def responder(prompt):
    try:
        inputs = tokenizer.encode(prompt, return_tensors="pt")
        outputs = model.generate(
            inputs,
            max_new_tokens=50,  # Limita tokens generados para evitar errores
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            num_return_sequences=1
        )
        respuesta = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return respuesta
    except Exception as e:
        return f"Error al generar la respuesta: {e}"

iface = gr.Interface(
    fn=responder,
    inputs="text",
    outputs="text",
    title="🤖 Bot GPT-2 Pequeño en Español",
    description="Bot de texto usando un modelo GPT-2 pequeño entrenado en español."
)

iface.launch()