import gradio as gr
import spaces
from huggingface_hub import hf_hub_download

import translator

REPO = 'ekvanox/mini-machine-translator'

model, token2idx, idx2token = translator.load(
    hf_hub_download(REPO, 'pico_model.pth'),
    hf_hub_download(REPO, 'pico_translator.vocab'))


# ZeroGPU is the only free hardware for Gradio Spaces and needs one @spaces.GPU function.
# The model is tiny, so it still runs on the CPU inside it.
@spaces.GPU(duration=10)
def translate(sentence):
    return translator.translate(model, token2idx, idx2token, sentence)


demo = gr.Interface(
    fn=translate,
    inputs=gr.Textbox(label='Swedish', lines=2),
    outputs=gr.Textbox(label='English'),
    title='Mini machine translator',
    description='A character-level encoder-decoder transformer (3+3 layers) '
                'trained on 20,532 Swedish-English Tatoeba pairs for EDAN20. '
                'It is small, so expect mistakes.',
    examples=['Kom och ät', 'Jag heter Erik', 'Var är stationen?'],
    flagging_mode='never')

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