import torch
import gradio as gr

# Load an En-De Transformer model trained on WMT'19 data:
en2de = torch.hub.load('pytorch/fairseq', 'transformer.wmt19.en-de.single_model', tokenizer='moses', bpe='fastbpe')
# Load an En-Fr Transformer model trained on WMT'14 data :
en2fr = torch.hub.load('pytorch/fairseq', 'transformer.wmt14.en-fr', tokenizer='moses', bpe='subword_nmt')

def translate(text, lang):
  if lang == "French":
    # Manually tokenize:
    en_toks = en2fr.tokenize(text)

    # Manually apply BPE:
    en_bpe = en2fr.apply_bpe(en_toks)

    # Manually binarize:
    en_bin = en2fr.binarize(en_bpe)

    # Generate five translations with top-k sampling:
    fr_bin = en2fr.generate(en_bin, beam=5, sampling=True, sampling_topk=20)

    # Convert one of the samples to a string and detokenize
    fr_sample = fr_bin[0]['tokens']
    fr_bpe = en2fr.string(fr_sample)
    fr_toks = en2fr.remove_bpe(fr_bpe)
    fr = en2fr.detokenize(fr_toks)
    return fr
  else: 
    # Translate from En-De
    de = en2de.translate(text)
    return de

inputs = [
  gr.inputs.Textbox(lines=5, label="Input Text in English"),
  gr.inputs.Radio(choices=["French", "German"], type="value", label="Output Language")
]
  
outputs = gr.outputs.Textbox(label="Output Text")

title = "Transformer (NMT)"
description = "Gradio demo for Transformer (NMT). To use it, simply add your text, or click one of the examples to load them. Read more at the links below."
article = """<p style='text-align: center'><a href='https://arxiv.org/abs/1806.00187'>Scaling Neural Machine Translation</a> | <a href='https://github.com/pytorch/fairseq/'>Github Repo</a></p>"""

examples = [
  ["Hello world!"],
  ["PyTorch Hub is a pre-trained model repository designed to facilitate research reproducibility."]
]

gr.Interface(translate, inputs, outputs, title=title, description=description, article=article, examples=examples, analytics_enabled=False).launch()