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)
#p = pipeline("automatic-speech-recognition", model = "facebook/wav2vec2-large-xlsr-53-spanish")


def transcribe(texto): # Desde un mensaje en español

    mensaje_voz = texto
    blob = tb(mensaje_voz)
    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
    
    return respuesta

def chatbot(input, history=[]):
    output = transcribe(input)
    history.append((input, output))
    return history, history
    
gr.Interface(
    fn=chatbot, 
    inputs=[
        "text",'state'
    ],
    outputs=[
        "chatbot",'state',
    ]
).launch()