tareesh commited on
Commit
2b81c45
·
1 Parent(s): eadff82

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +42 -0
app.py ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
3
+ from transformers import AutoModelForMaskedLM, AutoTokenizer
4
+
5
+ # Load the models and tokenizers
6
+ model_translation = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-one-to-many-mmt")
7
+ model_masked_lm = AutoModelForMaskedLM.from_pretrained("alabnii/jmedroberta-base-sentencepiece")
8
+ model_translation.eval()
9
+
10
+ tokenizer_translation = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-one-to-many-mmt", src_lang="en_XX")
11
+ tokenizer_masked_lm = AutoTokenizer.from_pretrained("alabnii/jmedroberta-base-sentencepiece")
12
+
13
+ text = st.text_area('Enter the text:')
14
+
15
+ if text:
16
+ model_inputs = tokenizer_translation(text, return_tensors="pt")
17
+ lg = st.text_input("Select Language: 1.Hindi, 2.Telugu, 3.Gujarati, 4.Bengali")
18
+ if lg==1:
19
+ generated_tokens = model_translation.generate(
20
+ **model_inputs,
21
+ forced_bos_token_id=tokenizer_translation.lang_code_to_id["hi_IN"]
22
+ )
23
+ else if lg==2:
24
+ generated_tokens = model_translation.generate(
25
+ **model_inputs,
26
+ forced_bos_token_id=tokenizer_translation.lang_code_to_id["te_IN"]
27
+ )
28
+ else if lg==3:
29
+ generated_tokens = model_translation.generate(
30
+ **model_inputs,
31
+ forced_bos_token_id=tokenizer_translation.lang_code_to_id["gu_IN"]
32
+ )
33
+ else if lg==4:
34
+ generated_tokens = model_translation.generate(
35
+ **model_inputs,
36
+ forced_bos_token_id=tokenizer_translation.lang_code_to_id["bn_IN"]
37
+ )
38
+ else:
39
+ return 0
40
+
41
+ translation = tokenizer_translation.batch_decode(generated_tokens, skip_special_tokens=True)
42
+ st.json(translation)