nevreal commited on
Commit
f9d944d
·
verified ·
1 Parent(s): 90deaee

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -56
app.py CHANGED
@@ -1,58 +1,3 @@
1
  import gradio as gr
2
- from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
3
 
4
- # Load the NLLB-200 model and tokenizer
5
- model_name = "facebook/nllb-200-distilled-600M"
6
- tokenizer = AutoTokenizer.from_pretrained(model_name)
7
- model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
8
-
9
- # Translation function
10
- def translate(text, src_lang, tgt_lang):
11
- # Convert the human-readable language names to their language codes
12
- src_lang_code = lang_dict[src_lang]
13
- tgt_lang_code = lang_dict[tgt_lang]
14
-
15
- # Tokenize input text with the source language code
16
- inputs = tokenizer(text, return_tensors="pt", src_lang=src_lang_code)
17
-
18
- # Generate translation with the target language code
19
- outputs = model.generate(**inputs, forced_bos_token_id=tokenizer.lang_code_to_id[tgt_lang_code])
20
-
21
- # Decode and return the translation
22
- translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
23
- return translation
24
-
25
- # List of languages supported by the model with their NLLB-200 codes
26
- lang_dict = {
27
- "English": "eng_Latn",
28
- "Spanish": "spa_Latn",
29
- "French": "fra_Latn",
30
- "German": "deu_Latn",
31
- "Chinese": "zho_Hans",
32
- "Hindi": "hin_Deva",
33
- "Indonesian": "ind_Latn", # Indonesian
34
- "Malaysian": "zsm_Latn", # Malaysian
35
- "Javanese": "jav_Latn", # Javanese
36
- "Sundanese": "sun_Latn", # Sundanese
37
- # Add more languages if needed
38
- }
39
-
40
- # Define the Gradio interface using Blocks
41
- with gr.Blocks() as ui:
42
- gr.Markdown("# Language Translation with NLLB-200 Model")
43
-
44
- with gr.Row():
45
- src_lang = gr.Dropdown(label="Source Language", choices=list(lang_dict.keys()), value="English")
46
- tgt_lang = gr.Dropdown(label="Target Language", choices=list(lang_dict.keys()), value="Indonesian")
47
-
48
- text_input = gr.Textbox(label="Enter text to translate", placeholder="Type your text here...")
49
- translate_button = gr.Button("Translate")
50
- translation_output = gr.Textbox(label="Translation Output")
51
-
52
- # Trigger translation on button click
53
- translate_button.click(fn=translate,
54
- inputs=[text_input, src_lang, tgt_lang],
55
- outputs=translation_output)
56
-
57
- # Launch the UI
58
- ui.launch()
 
1
  import gradio as gr
 
2
 
3
+ gr.load("models/facebook/nllb-200-distilled-600M").launch()