Spaces:
Running
Running
jhj0517
commited on
Commit
·
63fc551
1
Parent(s):
eeb932b
Add translation test
Browse files- tests/test_translation.py +52 -0
tests/test_translation.py
ADDED
@@ -0,0 +1,52 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from modules.translation.deepl_api import DeepLAPI
|
2 |
+
from modules.translation.nllb_inference import NLLBInference
|
3 |
+
from test_config import *
|
4 |
+
|
5 |
+
import os
|
6 |
+
import pytest
|
7 |
+
|
8 |
+
|
9 |
+
@pytest.mark.parametrize("model_size, file_path", [
|
10 |
+
(TEST_NLLB_MODEL, TEST_SUBTITLE_SRT_PATH),
|
11 |
+
(TEST_NLLB_MODEL, TEST_SUBTITLE_VTT_PATH),
|
12 |
+
])
|
13 |
+
def test_nllb_inference(
|
14 |
+
model_size: str,
|
15 |
+
file_path: str
|
16 |
+
):
|
17 |
+
nllb_inferencer = NLLBInference()
|
18 |
+
print("Device : ", nllb_inferencer.device)
|
19 |
+
|
20 |
+
result_str, file_paths = nllb_inferencer.translate_file(
|
21 |
+
fileobjs=[file_path],
|
22 |
+
model_size=model_size,
|
23 |
+
src_lang="eng_Latn",
|
24 |
+
tgt_lang="kor_Hang",
|
25 |
+
)
|
26 |
+
|
27 |
+
assert isinstance(result_str, str)
|
28 |
+
assert isinstance(file_paths[0], str)
|
29 |
+
|
30 |
+
|
31 |
+
@pytest.mark.parametrize("file_path", [
|
32 |
+
TEST_SUBTITLE_SRT_PATH,
|
33 |
+
TEST_SUBTITLE_VTT_PATH,
|
34 |
+
])
|
35 |
+
def test_deepl_api(
|
36 |
+
file_path: str
|
37 |
+
):
|
38 |
+
deepl_api = DeepLAPI()
|
39 |
+
|
40 |
+
api_key = os.getenv("DEEPL_API_KEY")
|
41 |
+
|
42 |
+
result_str, file_paths = deepl_api.translate_deepl(
|
43 |
+
auth_key=api_key,
|
44 |
+
fileobjs=[file_path],
|
45 |
+
source_lang="English",
|
46 |
+
target_lang="Korean",
|
47 |
+
is_pro=False,
|
48 |
+
add_timestamp=True,
|
49 |
+
)
|
50 |
+
|
51 |
+
assert isinstance(result_str, str)
|
52 |
+
assert isinstance(file_paths[0], str)
|