Update app.py
Browse files
app.py
CHANGED
@@ -1,19 +1,20 @@
|
|
1 |
import gradio as gr
|
2 |
from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
|
|
|
3 |
model = MBartForConditionalGeneration.from_pretrained("SnypzZz/Llama2-13b-Language-translate")
|
4 |
tokenizer = MBart50TokenizerFast.from_pretrained("SnypzZz/Llama2-13b-Language-translate", src_lang="en_XX")
|
5 |
|
6 |
-
dropdown = gr.Dropdown(["de_DE", "es_XX"], label="Choose Output Language")
|
7 |
|
8 |
-
def
|
9 |
model_inputs = tokenizer(input, return_tensors="pt")
|
10 |
generated_tokens = model.generate(
|
11 |
**model_inputs,
|
12 |
forced_bos_token_id=tokenizer.lang_code_to_id[dropdown_value]
|
13 |
)
|
14 |
-
|
15 |
-
|
16 |
-
return
|
17 |
|
18 |
-
iface = gr.Interface(fn=
|
19 |
iface.launch()
|
|
|
1 |
import gradio as gr
|
2 |
from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
|
3 |
+
|
4 |
model = MBartForConditionalGeneration.from_pretrained("SnypzZz/Llama2-13b-Language-translate")
|
5 |
tokenizer = MBart50TokenizerFast.from_pretrained("SnypzZz/Llama2-13b-Language-translate", src_lang="en_XX")
|
6 |
|
7 |
+
dropdown = gr.Dropdown(["de_DE", "es_XX", "fr_XX", "sv_SE", "ru_RU"], label="Choose Output Language")
|
8 |
|
9 |
+
def execute(input, dropdown_value):
|
10 |
model_inputs = tokenizer(input, return_tensors="pt")
|
11 |
generated_tokens = model.generate(
|
12 |
**model_inputs,
|
13 |
forced_bos_token_id=tokenizer.lang_code_to_id[dropdown_value]
|
14 |
)
|
15 |
+
output = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
|
16 |
+
output = output.strip("[]' ")
|
17 |
+
return output
|
18 |
|
19 |
+
iface = gr.Interface(fn=execute, inputs=["textbox", dropdown], outputs="textbox")
|
20 |
iface.launch()
|