jhj0517 commited on
Commit
63fc551
·
1 Parent(s): eeb932b

Add translation test

Browse files
Files changed (1) hide show
  1. 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)