|
import streamlit as st |
|
from transformers import MBartForConditionalGeneration, MBart50TokenizerFast |
|
from transformers import AutoModelForMaskedLM, AutoTokenizer |
|
|
|
|
|
model_translation = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-one-to-many-mmt") |
|
model_masked_lm = AutoModelForMaskedLM.from_pretrained("alabnii/jmedroberta-base-sentencepiece") |
|
model_translation.eval() |
|
|
|
tokenizer_translation = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-one-to-many-mmt", src_lang="en_XX") |
|
tokenizer_masked_lm = AutoTokenizer.from_pretrained("alabnii/jmedroberta-base-sentencepiece") |
|
|
|
text = st.text_area('Enter the text:') |
|
|
|
if text: |
|
model_inputs = tokenizer_translation(text, return_tensors="pt") |
|
lg = st.text_input("Select Language: hi.Hindi, te.Telugu, gu.Gujarati, bn.Bengali") |
|
if lg=='hi': |
|
generated_tokens = model_translation.generate( |
|
**model_inputs, |
|
forced_bos_token_id=tokenizer_translation.lang_code_to_id["hi_IN"] |
|
) |
|
translation = tokenizer_translation.batch_decode(generated_tokens, skip_special_tokens=True) |
|
st.write(translation) |
|
elif lg=='te': |
|
generated_tokens = model_translation.generate( |
|
**model_inputs, |
|
forced_bos_token_id=tokenizer_translation.lang_code_to_id["te_IN"] |
|
) |
|
translation = tokenizer_translation.batch_decode(generated_tokens, skip_special_tokens=True) |
|
st.write(translation) |
|
elif lg=='gu': |
|
generated_tokens = model_translation.generate( |
|
**model_inputs, |
|
forced_bos_token_id=tokenizer_translation.lang_code_to_id["gu_IN"] |
|
) |
|
translation = tokenizer_translation.batch_decode(generated_tokens, skip_special_tokens=True) |
|
st.write(translation) |
|
elif lg=='bn': |
|
generated_tokens = model_translation.generate( |
|
**model_inputs, |
|
forced_bos_token_id=tokenizer_translation.lang_code_to_id["bn_IN"] |
|
) |
|
translation = tokenizer_translation.batch_decode(generated_tokens, skip_special_tokens=True) |
|
st.json(translation) |
|
else: |
|
st.write('invalid choice!') |
