File size: 771 Bytes
3e2e722
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
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()