Spaces:
Running
on
Zero
Running
on
Zero
Commit
•
3f23d73
1
Parent(s):
6bcde50
Refactor code to improve performance and readability
Browse files
app.py
CHANGED
@@ -14,11 +14,17 @@ flores_codes = list(code_mapping.keys())
|
|
14 |
|
15 |
def load_model():
|
16 |
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME).to(device)
|
17 |
-
|
18 |
-
return model, tokenizer
|
19 |
|
20 |
|
21 |
-
model
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
22 |
|
23 |
|
24 |
@spaces.GPU
|
@@ -29,12 +35,10 @@ def translate(
|
|
29 |
window_size: int = 800,
|
30 |
overlap_size: int = 200,
|
31 |
):
|
|
|
|
|
32 |
input_tokens = (
|
33 |
-
tokenizer(text, return_tensors="pt"
|
34 |
-
.input_ids[0]
|
35 |
-
.cpu()
|
36 |
-
.numpy()
|
37 |
-
.tolist()
|
38 |
)
|
39 |
translated_chunks = []
|
40 |
|
@@ -51,8 +55,7 @@ def translate(
|
|
51 |
)
|
52 |
translated_chunks.append(translated_chunk)
|
53 |
|
54 |
-
|
55 |
-
return translated_text
|
56 |
|
57 |
|
58 |
description = """
|
|
|
14 |
|
15 |
def load_model():
|
16 |
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME).to(device)
|
17 |
+
return model
|
|
|
18 |
|
19 |
|
20 |
+
model = load_model()
|
21 |
+
|
22 |
+
|
23 |
+
def load_tokenizer(src_lang, tgt_lang):
|
24 |
+
tokenizer = AutoTokenizer.from_pretrained(
|
25 |
+
MODEL_NAME, src_lang=code_mapping[src_lang], tgt_lang=code_mapping[tgt_lang]
|
26 |
+
)
|
27 |
+
return tokenizer
|
28 |
|
29 |
|
30 |
@spaces.GPU
|
|
|
35 |
window_size: int = 800,
|
36 |
overlap_size: int = 200,
|
37 |
):
|
38 |
+
tokenizer = load_tokenizer(src_lang, tgt_lang)
|
39 |
+
|
40 |
input_tokens = (
|
41 |
+
tokenizer(text, return_tensors="pt").input_ids[0].cpu().numpy().tolist()
|
|
|
|
|
|
|
|
|
42 |
)
|
43 |
translated_chunks = []
|
44 |
|
|
|
55 |
)
|
56 |
translated_chunks.append(translated_chunk)
|
57 |
|
58 |
+
return " ".join(translated_chunks)
|
|
|
59 |
|
60 |
|
61 |
description = """
|