import gradio as gr
from transformers import AutoTokenizer
from optimum.onnxruntime import ORTModelForSeq2SeqLM
from translator import Translator
from translation_pipeline import TranslationPipeline

tokenizer = AutoTokenizer.from_pretrained("hon9kon9ize/bart-translation-zh-yue-onnx")
model = ORTModelForSeq2SeqLM.from_pretrained(
    "hon9kon9ize/bart-translation-zh-yue-onnx",
    provider="CPUExecutionProvider",
    encoder_file_name="encoder_model_quantized.onnx",
    decoder_file_name="decoder_model_quantized.onnx",
    decoder_file_with_past_name="decoder_with_past_model_quantized.onnx",
)
pipe = TranslationPipeline(model, tokenizer)
translator = Translator(pipe, max_length=512, batch_size=1)


def demo_process(input_text):
    return translator(input_text)[0]


demo = gr.Interface(
    fn=demo_process,
    inputs=gr.Textbox(lines=5, label="Input Text"),
    outputs="text",
    title=f"Chinese to Cantonese Translator",
    description="This is a demo of the Chinese to Cantonese Translator.",
    examples=[
        [
            "近年成为许多港人热门移居地的英国中部城巿诺定咸（又译诺丁汉，Nottingham），多年来一直面对财政困境，市议会周三（11月29日）宣布破产，是继英国第二大城市伯明翰今年9月宣布破产后，近期「爆煲」的另一个英国主要城市。诺定咸除了维持法例规定必须提供的服务外，巿政府将暂停所有非必要的公共开支。"
        ]
    ],
    cache_examples=True,
)
demo.launch(server_name="0.0.0.0")