import gradio as gr
from transformers import AutoTokenizer, GPTJForCausalLM

# Cargar el modelo en español
model_name = "mrm8488/bertin-gpt-j-6B-ES-v1-8bit"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = GPTJForCausalLM.from_pretrained(model_name)

# Definir la función de respuesta para el chatbot
def chatbot(message):
    inputs = tokenizer(message, return_tensors="pt")
    outputs = model.generate(inputs.input_ids, max_length=150)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# Crear la interfaz de Gradio
demo = gr.Interface(
    fn=chatbot, 
    inputs="text", 
    outputs="text", 
    title="Chatbot en Español"
)

# Ejecutar la aplicación
if __name__ == "__main__":
    demo.launch()