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