Update app.py
Browse files
app.py
CHANGED
@@ -68,12 +68,12 @@ dict_map = {
|
|
68 |
}
|
69 |
|
70 |
tokenizer_vi2en = AutoTokenizer.from_pretrained("vinai/vinai-translate-vi2en-v2", src_lang="vi_VN")
|
71 |
-
model_vi2en = AutoModelForSeq2SeqLM.from_pretrained("vinai/vinai-translate-vi2en-v2")
|
72 |
|
73 |
def translate_vi2en(vi_text: str) -> str:
|
74 |
for i, j in dict_map.items():
|
75 |
vi_text = vi_text.replace(i, j)
|
76 |
-
input_ids = tokenizer_vi2en(vi_text, return_tensors="pt").input_ids
|
77 |
output_ids = model_vi2en.generate(
|
78 |
input_ids,
|
79 |
decoder_start_token_id=tokenizer_vi2en.lang_code_to_id["en_XX"],
|
|
|
68 |
}
|
69 |
|
70 |
tokenizer_vi2en = AutoTokenizer.from_pretrained("vinai/vinai-translate-vi2en-v2", src_lang="vi_VN")
|
71 |
+
model_vi2en = AutoModelForSeq2SeqLM.from_pretrained("vinai/vinai-translate-vi2en-v2",device_map="auto")
|
72 |
|
73 |
def translate_vi2en(vi_text: str) -> str:
|
74 |
for i, j in dict_map.items():
|
75 |
vi_text = vi_text.replace(i, j)
|
76 |
+
input_ids = tokenizer_vi2en(vi_text, return_tensors="pt")to("cuda").input_ids
|
77 |
output_ids = model_vi2en.generate(
|
78 |
input_ids,
|
79 |
decoder_start_token_id=tokenizer_vi2en.lang_code_to_id["en_XX"],
|