from transformers import (
    AutoTokenizer,
    AutoModelForSeq2SeqLM,
    pipeline
)
from textblob import TextBlob as tb
import gradio as gr


tokenizer = AutoTokenizer.from_pretrained("microsoft/GODEL-v1_1-base-seq2seq")
model = AutoModelForSeq2SeqLM.from_pretrained("microsoft/GODEL-v1_1-base-seq2seq")
pipe = pipeline(model="aware-ai/whisper-base-german")


def translate(text):
    blob = tb(text)
    translation = str(blob.translate(from_lang='de',to='en'))
    return translation


def translate_to_de(text):
    blob = tb(text)
    translation = str(blob.translate(from_lang='en',to='de'))
    return translation
    
    
def transcribe(audio):
    text = pipe(audio)["text"]
    return text


def generate(input, knowledge):

    if knowledge == '':
        pass
    else:
        knowledge = translate(knowledge)
        
    input = translate(input)

    top_p = 1
    min_length = 8
    max_length = 64

    instruction = 'given a dialog context and related knowledge, you need to answer the question based on the knowledge.'

    if knowledge != '':
        knowledge = '[KNOWLEDGE] ' + knowledge

    dialog = ' EOS '.join([input])
    query = f"{instruction} [CONTEXT] {dialog} {knowledge}"

    input_ids = tokenizer(f"{query}", return_tensors="pt").input_ids
    outputs = model.generate(input_ids, min_length=int(
        min_length), max_length=int(max_length), top_p=top_p, do_sample=True)
    output = tokenizer.decode(outputs[0], skip_special_tokens=True)

    output = translate_to_de(output)
    return output

with gr.Blocks() as app:

    conocimiento = gr.Textbox(label='Conocimiento',lines=7,max_lines=7)

    voice = gr.Audio(source='microphone',type='filepath')
    
    send_button = gr.Button(value='Transcribir')
    transc = gr.Textbox(label='Transcripción',value='',)
    
    button2 = gr.Button(value='Respuesta de la IA')
    respuesta = gr.Textbox(label='Respuesta',interactive=False,value='')

    send_button.click(fn=transcribe,inputs=voice,outputs=transc)
    button2.click(fn=generate,inputs=[transc,conocimiento],outputs=respuesta)

app.launch()