zedzek / app.py
anzorq's picture
add app.py
3e2e722
raw
history blame
771 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
iface = gr.Interface(fn=translate, inputs="text", outputs="text")
iface.launch()