| | 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] |
