import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")
infer = AutoModelForSeq2SeqLM.from_pretrained("Qilex/bart-largeEN-ME")

def translate(sentence):
    input_ids = tokenizer(sentence, return_tensors="pt").input_ids
    outputs = infer.generate(input_ids, max_new_tokens = len(sentence.split(' '))*10)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

title = "English to Middle English Translator"
description = """
This translator is trained on about 80,000 English/Middle English paired sentences. 
<br>
It's still a work in progress.
"""
article = ''''''
gr.Interface(
    fn=translate,
    inputs=gr.Textbox(lines=1, placeholder="Enter text to translate."),
    outputs="text",
    title=title,
    description=description,
    article = article,
).launch()