Kleber commited on
Commit
e5f115a
1 Parent(s): fd62163

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +58 -0
app.py ADDED
@@ -0,0 +1,58 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
3
+ import torch
4
+
5
+ LANGS = ["kin_Latn","eng_Latn"]
6
+ TASK = "translation"
7
+ # CKPT = "DigitalUmuganda/Finetuned-NLLB"
8
+ # MODELS = ["facebook/nllb-200-distilled-600M","DigitalUmuganda/Finetuned-NLLB"]
9
+ # model = AutoModelForSeq2SeqLM.from_pretrained(CKPT)
10
+ # tokenizer = AutoTokenizer.from_pretrained(CKPT)
11
+
12
+ device = 0 if torch.cuda.is_available() else -1
13
+
14
+ #general_model = AutoModelForSeq2SeqLM.from_pretrained("mbazaNLP/Nllb_finetuned_general_en_kin")
15
+ education_model = AutoModelForSeq2SeqLM.from_pretrained("mbazaNLP/Nllb_finetuned_education_en_kin")
16
+ #tourism_model = AutoModelForSeq2SeqLM.from_pretrained("mbazaNLP/Nllb_finetuned_tourism_en_kin")
17
+ #MODELS = {"General model":general_model_model,"Education model":education_model,"Tourism model":tourism_model}
18
+ #MODELS = {"Education model":education_model,"Tourism model":tourism_model}
19
+
20
+
21
+ tokenizer = AutoTokenizer.from_pretrained("mbazaNLP/Nllb_finetuned_general_en_kin")
22
+ # def translate(text, src_lang, tgt_lang, max_length=400):
23
+ def translate(CKPT,text, src_lang, tgt_lang, max_length=400):
24
+
25
+ """
26
+ Translate the text from source lang to target lang
27
+ """
28
+
29
+
30
+ translation_pipeline = pipeline(TASK,
31
+ model=education_model,
32
+ tokenizer=tokenizer,
33
+ src_lang=src_lang,
34
+ tgt_lang=tgt_lang,
35
+ max_length=max_length,
36
+ device=device)
37
+
38
+ result = translation_pipeline(text)
39
+ return result[0]['translation_text']
40
+
41
+
42
+ gr.Interface(
43
+ translate,
44
+ [
45
+ #gr.components.Dropdown(label="choose a model",choices=MODELS),
46
+ gr.components.Textbox(label="Text"),
47
+ gr.components.Dropdown(label="Source Language", choices=LANGS),
48
+ gr.components.Dropdown(label="Target Language", choices=LANGS),
49
+ #gr.components.Slider(8, 512, value=400, step=8, label="Max Length")
50
+ ],
51
+ ["text"],
52
+ #examples=examples,
53
+ # article=article,
54
+ cache_examples=False,
55
+ title="Finetuned-NLLB-EN-KIN",
56
+ #description=description
57
+ ).launch()
58
+