import gradio as gr
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

def Auto_Translate(sentences):
  inputs = tokenizer(sentences, return_tensors="pt")
  translated_tokens = model.generate(**inputs, forced_bos_token_id=tokenizer.lang_code_to_id["eng_Latn"], max_length=360)
  return tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]

access_token = "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"

tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M", use_auth_token= access_token , src_lang="chn_Latn")
model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M", use_auth_token=access_token)
demo = gr.Interface(
    fn=Auto_Translate,
    inputs=gr.Textbox(lines=2, placeholder="Input a Chinese sentence!"),
    outputs="text", 
    description="Chinese-English Neural MT based on Meta's NLLB200"
)
demo.launch()