from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from transformers import pipeline
import gradio as gr



model_name = "Helsinki-NLP/opus-mt-en-zh"
tokenizer = AutoTokenizer.from_pretrained(model_name)
# model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
pipe = pipeline("translation", model=model_name, tokenizer=tokenizer)





def getTranslateResult(inputTextValue):
    testRes = pipe(inputTextValue)
    return testRes[0]["translation_text"]




with gr.Blocks() as translateDemo:
    with gr.Row():
        with gr.Column():
            inputText = gr.Textbox(label="Please enter the translation content, only Chinese is supported...")
            submitBtn = gr.Button(variant="primary", value="Submit")
            outText = gr.Textbox(label='Translation results are shown here...')


            inputText.submit(
                fn=getTranslateResult,
                inputs=inputText,
                outputs=outText
            )

            submitBtn.click(
                fn=getTranslateResult,
                inputs=inputText,
                outputs=outText
            )

translateDemo.queue().launch()