import gradio as gr
import transformers

encoder_model_name = "cl-tohoku/bert-base-japanese-v2"
decoder_model_name = "skt/kogpt2-base-v2"
src_tokenizer = transformers.BertJapaneseTokenizer.from_pretrained(encoder_model_name)
trg_tokenizer = transformers.PreTrainedTokenizerFast.from_pretrained(decoder_model_name)
model = transformers.EncoderDecoderModel.from_pretrained("sappho192/ffxiv-ja-ko-translator")


def translate(text_src):
    embeddings = src_tokenizer(text_src, return_attention_mask=False, return_token_type_ids=False, return_tensors='pt')
    embeddings = {k: v for k, v in embeddings.items()}
    output = model.generate(**embeddings, max_length=500)[0, 1:-1]
    text_trg = trg_tokenizer.decode(output.cpu())
    return text_trg


def endpoint(sentence):
    return translate(sentence)


# demo = gr.Interface(fn=endpoint, inputs="text", outputs="text")
with gr.Blocks() as demo:
    input = gr.Textbox(label="Sentence")
    output = gr.Textbox(label="Result")
    btn = gr.Button(value="Submit")
    btn.click(endpoint, inputs=[input], outputs=[output])
    
    gr.Markdown("## Examples")
    gr.Markdown(
    """
    For now, the model can translate the words included in Auto-Translate dictionary.
    현재 버전의 번역기는 상용구에 포함된 단어나 표현이 들어간 문장을 어느정도 번역할 수 있습니다.
    학습에 사용된 데이터 및 번역 데이터는 글로벌 서버의 일본어 클라이언트에서만 취득하였으며, 한국 서비스에서 제공되는 번역 데이터를 일절 사용하지 않았기 때문에 용어의 번역 결과가 대부분 직역의 형태로 제공됩니다.
    """)
    gr.Examples(
        [["ギルガメッシュ討伐戦に行ってきます。一緒に行きましょうか？"], 
         ["絶アルテマウェポン破壊作戦をクリアし事ありますか？"],
         ["美容師の呼び鈴を使って髪方を変えますよ。"]],
        [input],
        output,
        endpoint,
        cache_examples=False
    )
    
if __name__ == "__main__":
    demo.launch()