zedzek / app.py
anzorq's picture
~ app.py
a15bcba
raw
history blame
801 Bytes
import gradio as gr
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
model_path = "anzorq/m2m100_418M_ft_ru-kbd_44K"
src_lang="ru"
tgt_lang="zu"
tokenizer = AutoTokenizer.from_pretrained(model_path, src_lang=src_lang)
model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
def translate(text):
inputs = tokenizer(text, return_tensors="pt")
translated_tokens = model.generate(
**inputs, forced_bos_token_id=tokenizer.lang_code_to_id[tgt_lang], num_beams=4, num_return_sequences=4
)
translations = []
for translation in tokenizer.batch_decode(translated_tokens, skip_special_tokens=True):
translations.append(translation)
return translations
output = gr.outputs.Textbox()
iface = gr.Interface(fn=translate, inputs="text", outputs=output)
iface.launch()