|
|
import re |
|
|
|
|
|
languge_settings = { |
|
|
'python': { |
|
|
'full_name': 'Python', |
|
|
'indent': 4, |
|
|
}, |
|
|
'cpp': { |
|
|
'full_name': 'cpp', |
|
|
'indent': 0, |
|
|
'main': "int main()", |
|
|
}, |
|
|
'java': { |
|
|
'full_name': 'Java', |
|
|
'indent': 4, |
|
|
'main': "public static void main", |
|
|
}, |
|
|
'cs': { |
|
|
'full_name': "csharp", |
|
|
'indent': 0, |
|
|
'main': "public static void Main", |
|
|
}, |
|
|
'php': { |
|
|
'full_name': "PHP", |
|
|
'indent': 0, |
|
|
}, |
|
|
'ts': { |
|
|
'full_name': "TypeScript", |
|
|
'indent': 0, |
|
|
}, |
|
|
'js': { |
|
|
'full_name': "JavaScript", |
|
|
'indent': 0 |
|
|
}, |
|
|
'sh': { |
|
|
'full_name': "Bash", |
|
|
'indent': 0 |
|
|
} |
|
|
} |
|
|
|
|
|
def get_function_name(question: str, lang: str): |
|
|
func_lines = [x for x in question.strip().split('\n') if x.strip()] |
|
|
|
|
|
if lang.lower() == 'python': |
|
|
func_idx = [i for i in range(len(func_lines)) if func_lines[i].startswith("def ")][-1] |
|
|
func_name = func_lines[func_idx].split('(')[0].strip() |
|
|
func_prefix = "\n".join(func_lines[:func_idx]) |
|
|
return func_name, func_prefix |
|
|
|
|
|
func_name = func_lines[-1].split('{')[0].strip() |
|
|
func_prefix = "\n".join(func_lines[:-1]) |
|
|
return func_name, func_prefix |
|
|
|
|
|
def extract_generation_code(example: str, lang_code: str, verbose: bool=False): |
|
|
task_id = example['task_id'] |
|
|
output = example.get('output', example.get("gpt_completion")) |
|
|
question = example["prompt"].strip() |
|
|
setting = languge_settings[lang_code] |
|
|
lang = setting['full_name'] |
|
|
indent = setting['indent'] |
|
|
|
|
|
try: |
|
|
code_block: str = re.findall(f'```{lang.lower()}\n(.*?)```', output, re.DOTALL | re.IGNORECASE)[0] |
|
|
if verbose: |
|
|
print(">>> Task: {}\n{}".format(task_id, code_block)) |
|
|
|
|
|
|
|
|
if setting.get('main', None) and setting['main'] in code_block: |
|
|
main_start = code_block.index(setting['main']) |
|
|
code_block = code_block[:main_start] |
|
|
|
|
|
func_name, func_prefix = get_function_name(question, lang) |
|
|
|
|
|
try: |
|
|
start = code_block.lower().index(func_name.lower()) |
|
|
indent = 0 |
|
|
while start - indent >= 0 and code_block[start - indent-1] == ' ': |
|
|
indent += 1 |
|
|
|
|
|
try: |
|
|
end = code_block.rindex('\n' + ' '*indent + '}') |
|
|
except: |
|
|
end = len(code_block) |
|
|
except: |
|
|
start = 0 |
|
|
try: |
|
|
end = code_block.rindex('\n' + ' '*indent + '}') |
|
|
except: |
|
|
end = len(code_block) |
|
|
|
|
|
body = code_block[start:end] |
|
|
|
|
|
if lang_code.lower() in ['php', 'ts', 'js']: |
|
|
body += '\n' + ' '*indent + '}' |
|
|
|
|
|
generation = func_prefix + '\n' + body + '\n' |
|
|
example['generation'] = generation |
|
|
|
|
|
except Exception as ex: |
|
|
print("Failed to extract code block with error `{}`:\n>>> Task: {}\n>>> Output:\n{}".format( |
|
|
ex, task_id, output |
|
|
)) |
|
|
example['generation'] = example['prompt'] + '\n' + output |
|
|
|
|
|
return example |
|
|
|
|
|
def cleanup_code( |
|
|
code: str, |
|
|
language_type: str = None, |
|
|
dataset: str = None, |
|
|
issft: bool = False, |
|
|
stop_words = [] |
|
|
): |
|
|
""" |
|
|
Cleans up the generated code. |
|
|
""" |
|
|
|
|
|
if language_type.lower() == "python": |
|
|
if issft: |
|
|
code = _clean_python_code_for_sft(code) |
|
|
stop_words = ["\ndef", "\nclass", "\nif", "\n#", "\nprint"] |
|
|
code = _truncate_code_at_stopwords(code, stop_words) |
|
|
elif language_type.lower() == "ts": |
|
|
code = _truncate_code_at_stopwords(code, stop_words + ["\nexport", "\nimport", "\nexport default", "\nimport default", "\nconsole.log"]) |
|
|
else: |
|
|
code = _truncate_code_at_stopwords(code, stop_words) |
|
|
|
|
|
return code |
|
|
|
|
|
def _clean_python_code_for_sft(code): |
|
|
code = code.replace("\r", "") |
|
|
if "```python" in code: |
|
|
code_start_idx = code.index("```python") |
|
|
code = code[code_start_idx:].replace("```python", "").strip() |
|
|
end_idx = code.find("```") if "```" in code else len(code) |
|
|
code = code[:end_idx].strip() |
|
|
|
|
|
return code |
|
|
|
|
|
def _truncate_code_at_stopwords(code, stop_words): |
|
|
min_stop_idx = len(code) |
|
|
for stop_word in stop_words: |
|
|
stop_index = code.find(stop_word) |
|
|
if 0 <= stop_index < min_stop_idx: |
|
|
min_stop_idx = stop_index |
|
|
return code[:min_stop_idx] |
