Babyloncoder commited on
Commit
0777539
1 Parent(s): 63d1a49

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +55 -0
app.py ADDED
@@ -0,0 +1,55 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import torch
3
+ import gradio as gr
4
+ import time
5
+ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
6
+ from code_flores_latest import flores_codes_latest
7
+
8
+
9
+ def load_models():
10
+ model_name_dict = {'nllb-distilled-600M': 'facebook/nllb-200-distilled-600M'}
11
+ model_dict = {}
12
+ for call_name, real_name in model_name_dict.items():
13
+ print('\tLoading model: %s' % call_name)
14
+ model = AutoModelForSeq2SeqLM.from_pretrained(real_name)
15
+ tokenizer = AutoTokenizer.from_pretrained(real_name)
16
+ model_dict[call_name + '_model'] = model
17
+ model_dict[call_name + '_tokenizer'] = tokenizer
18
+ return model_dict
19
+
20
+
21
+ def translation(source, target, text):
22
+ model_name = 'nllb-distilled-600M'
23
+ source_code = flores_codes_latest.get(source, None)
24
+ target_code = flores_codes_latest.get(target, None)
25
+ if not source_code or not target_code:
26
+ return "<p>Error: Language code not found.</p>"
27
+
28
+ model = model_dict[model_name + '_model']
29
+ tokenizer = model_dict[model_name + '_tokenizer']
30
+ translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=source_code, tgt_lang=target_code)
31
+ output = translator(text, max_length=400)
32
+ output_text = output[0]['translation_text']
33
+
34
+ formatted_output = f"<p><strong>Original Text ({source}):</strong><br>{text}</p><p><strong>Translated Text ({target}):</strong> <span style='color: red;'>{output_text}</span></p>"
35
+ return formatted_output
36
+
37
+
38
+ if __name__ == '__main__':
39
+ print('\tInitializing models')
40
+ model_dict = load_models()
41
+ lang_names = list(flores_codes_latest.keys())
42
+ source_dropdown = gr.Dropdown(lang_names, label='Source', allow_custom_value=True)
43
+ target_dropdown = gr.Dropdown(lang_names, label='Target', allow_custom_value=True)
44
+ textbox = gr.Textbox(lines=5, label="Input text")
45
+ title = "nllb-distilled-600M - Example implementation"
46
+
47
+ description = f"Nots: please note that not all translations are accurate, and some models codes are not accepted . "
48
+
49
+ initial_output_value = "<p>Translation results will appear here.</p>"
50
+ output_html = gr.HTML(label="Translation Result", value=initial_output_value)
51
+
52
+ iface = gr.Interface(fn=translation, inputs=[source_dropdown, target_dropdown, textbox], outputs=output_html,
53
+ title=title, description=description)
54
+ iface.launch()
55
+