import gradio as gr
from transformers import M2M100ForConditionalGeneration
from tokenization_small100 import SMALL100Tokenizer

langs = """af,am,ar,ast,az,ba,be,bg,bn,br,bs,ca,ceb,cs,cy,da,de,el,en,es,et,fa,ff,fi,fr,fy,ga,gd,gl,gu,ha,he,hi,hr,ht,hu,hy,id,ig,ilo,is,it,ja,jv,ka,kk,km,kn,ko,lb,lg,ln,lo,lt,lv,mg,mk,ml,mn,mr,ms,my,ne,nl,no,ns,oc,or,pa,pl,ps,pt,ro,ru,sd,si,sk,sl,so,sq,sr,ss,su,sv,sw,ta,th,tl,tn,tr,uk,ur,uz,vi,wo,xh,yi,yo,zh,zu"""
lang_list = langs.split(',')

model = M2M100ForConditionalGeneration.from_pretrained("alirezamsh/small100")
tokenizer = SMALL100Tokenizer.from_pretrained("alirezamsh/small100")

def translate(lang, text):
    tokenizer.tgt_lang = lang
    encoded_text = tokenizer(text, return_tensors="pt")
    generated_tokens = model.generate(**encoded_text)
    return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]

with gr.Blocks(analytics_enabled=False) as app:
    
    Source = gr.Textbox( label="Source" )
    Language = gr.Dropdown( lang_list, label="Language" )
    Translate = gr.Button( "Translate" )
    Result = gr.Textbox( label="Result" )
    Info = gr.Markdown( "# [$hyoo_lingua](https://lingua.hyoo.ru/)" )

    Translate.click(
        translate,
        inputs=[ Language, Source ],
        outputs=[Result],
        api_name="translate",
    )

    app.launch( inline=True )
    block.queue( concurrency_count=2 )