import gradio as gr
from transformers import MBartForConditionalGeneration, MBart50Tokenizer

def download_model():
    model_name = "facebook/mbart-large-50-many-to-many-mmt"
    model = MBartForConditionalGeneration.from_pretrained(model_name)
    tokenizer = MBart50Tokenizer.from_pretrained(model_name)
    return model, tokenizer

model, tokenizer = download_model()

def Translation_demo(input_text):
    model_inputs = tokenizer(input_text, return_tensors="pt")
    generated_tokens = model.generate(**model_inputs,forced_bos_token_id=tokenizer.lang_code_to_id["en_XX"])
    translation = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
    return translation

title ='Translate the sentences'
description = "Language Translation"
examples = [
    ["I am proud of being an Indian"],
    ["I am happy with the performance of Indian Hockey team"],
    ["We are all born with a divine fire in us. Our efforts should be to give wings to this fire”. President Kalam’s inspiring autobiography has touched the hearts of people all over the world. The book talks about the his humble beginnings in a small village in Tamil Nadu and his ascent to one of the most powerful positions in the country."]
]

translation_demo= gr.Interface(fn=Translation_demo, inputs = 'text', outputs='text', title = title, description = description ,examples = examples)

#translation_demo.launch(inline = False)