import gradio as gr
import json
from nameder import init_model_ner, get_entity_results
from speech2text import init_model_trans, transcribe
from translation import translate
from resources import NER_Response, NER_Request, entity_labels_sample, set_start, audit_elapsedtime


def translation_to_english(text: str):
   resultado = translate(text)
   return resultado

def transcription(audio: bytes):
    
    s2t = init_model_trans()
    return transcribe(audio, s2t)

def named_entity_recognition(text: str):
    tokenizer, ner = init_model_ner()
    # print('NER:',ner)
    result = get_entity_results(entities_list=entity_labels_sample, 
                                model=ner,
                                tokenizer=tokenizer,
                                text=text)
    print('result:',result,type(result))
    return result

def get_lead(audio: bytes):
    start = set_start()
    transcribe = transcription(audio)
    translate = translation_to_english(transcribe)
    ner = named_entity_recognition(NER_Request(
        entities=entity_labels_sample,
        text=translate
    ))
    audit_elapsedtime("VoiceLead", start)
    return ner

audio_input = gr.Microphone(
    label="Record your audio"
)
text_output = gr.Textbox(
            label="Labels",
            info="",
            lines=9,
            value=""
        )
demo = gr.Interface(
    fn=named_entity_recognition,
    description= "Get the ",
    inputs=[audio_input],
    outputs=[text_output],
    title="VoiceLead"
)

if __name__ == "__main__":
    demo.launch()