import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

model_translate = AutoModelForSeq2SeqLM.from_pretrained("h4d35/Translator")
tokenizer_translate = AutoTokenizer.from_pretrained("h4d35/Translator")
inlang='hi'
outlang='en'
tokenizer_translate.src_lang = inlang

def translate(text):    
    encoded_hi = tokenizer_translate(text, return_tensors="pt")
    generated_tokens = model_translate.generate(**encoded_hi, forced_bos_token_id=tokenizer_translate.get_lang_id(outlang))
    return tokenizer_translate.batch_decode(generated_tokens, skip_special_tokens=True)[0]
    
interface = gr.Interface(
    fn = translate,
    inputs = 'text',
    outputs = 'text'
)

interface.launch()