from transformers import BlenderbotTokenizer, BlenderbotForConditionalGeneration, pipeline
import gradio as gr
from textblob import TextBlob as tb

mname = "facebook/blenderbot-400M-distill"
model = BlenderbotForConditionalGeneration.from_pretrained(mname)
tokenizer = BlenderbotTokenizer.from_pretrained(mname)


def chat(mensaje, history): # Desde un mensaje en español

    history = history or []

    blob = tb(mensaje)
    MESSAGE = str(blob.translate(from_lang='es', to='en')) # Traduce el mensaje al inglés

    inputs = tokenizer(MESSAGE, return_tensors="pt") # Tokeniza el mensaje traducido
    reply_ids = model.generate(**inputs)

    response = tokenizer.batch_decode(reply_ids)[0] # Genera la respuesta (en inglés)

    blob_2 = tb(response)
    respuesta = str(blob_2.translate(from_lang='en', to='es'))[4:-5] # Traduce la respuesta al español

    history.append((mensaje, respuesta))
    
    return history, history

chatbot = gr.Chatbot().style(color_map=("red", "blue"))


demo = gr.Interface(
    chat,
    ["text", "state"],
    [chatbot, "state"],
    allow_flagging="never",
)

demo.launch(debug=True)