-
Notifications
You must be signed in to change notification settings - Fork 0
/
translate.py
41 lines (30 loc) · 1.11 KB
/
translate.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
from utils.utils import read_problems, write_jsonl
import deepl
import os
def translate(prompt, translator):
# extract docstring
doc_start = prompt.find('"""')
doc_end = prompt.find('"""', doc_start + 3)
if doc_start == -1 or doc_end == -1:
doc_start = prompt.find("'''")
doc_end = prompt.find("'''", doc_start + 3)
docstring = prompt[doc_start + 3: doc_end]
# translate
translated_docstring = translator.translate_text(docstring, target_lang="KO").text
new_prompt = prompt[:doc_start + 3] + translated_docstring + prompt[doc_end:]
# print('original prompt:\n')
# print(prompt)
# print('new_prompt:\n')
# print(new_prompt)
return new_prompt
api_key, save_path = None, None # assign!
translator = deepl.Translator(api_key)
dataset = read_problems('HumanEval')
for pid in dataset:
result = []
example = dataset[pid]
prompt = example['prompt']
translated_prompt = translate(prompt, translator)
example['translated_prompt'] = translated_prompt
result.append(example)
write_jsonl(save_path, result, append=True)