parse code blocks

alfworld_runs/utils.py

@@ -8,7 +8,7 @@ from tenacity import (
 )
 
 from typing import Optional, List
-if sys.version_info >= (3, 7):
+if sys.version_info >= (3, 8):
     from typing import Literal
 else:
     from typing_extensions import Literal

programming_runs/generators/generator_utils.py

@@ -1,6 +1,6 @@
 from generators.model import ModelBase, Message
 import random
-import re
+from parse import parse_code_block, add_code_block
 
 from typing import Union, List, Optional, Callable
 
@@ -8,16 +8,6 @@ from typing import Union, List, Optional, Callable
 
 USE_PYTHON_CODEBLOCK_INSTRUCTION = "Use a Python code block to write your response. For example:\n```python\nprint('Hello world!')\n```"
 
-def parse_python_code(string: str) -> Optional[str]:
-    code_pattern = r"```python\n(.*?)\n```"
-    match = re.search(code_pattern, string, re.DOTALL)
-    if match:
-        return match.group(1)
-    else:
-        return None
-
-def add_code_block(string: str) -> str:
-    return f"```python\n{string}\n```"
 
 def generic_generate_func_impl(
     func_sig: str,
@@ -33,7 +23,8 @@ def generic_generate_func_impl(
     simple_chat_instruction: str,
     reflexion_completion_instruction: str,
     simple_completion_instruction: str,
-    fix_body: Callable[[str], str]
+    parse_code_block: Callable[[str], str],
+    add_code_block: Callable[[str], str]
 ) -> Union[str, List[str]]:
     if strategy != "reflexion" and strategy != "simple":
         raise ValueError(
@@ -101,12 +92,12 @@ def generic_generate_func_impl(
 
     if num_comps == 1:
         assert isinstance(func_bodies, str)
-        func_body_str = parse_python_code(func_bodies)
+        func_body_str = parse_code_block(func_bodies)
         print_generated_func_body(func_body_str)
         return func_body_str
 
     else:
-        func_bodies = [parse_python_code(func_body) for func_body in func_bodies]
+        func_bodies = [parse_code_block(func_body) for func_body in func_bodies]
         print_generated_func_body("\n\n".join(func_bodies))
         return func_bodies
 

programming_runs/generators/parse_indent.py

@@ -1,66 +0,0 @@
-test0 = """    res = 0
-    for i in range(len(l)):
-        res += l[i] * 2 ** i
-    return res"""
-
-test1 = """res = 0
-    for i in range(len(l)):
-        res += l[i] * 2 ** i
-    return res"""
-
-test2 = """res = 0
-for i in range(len(l)):
-    res += l[i] * 2 ** i
-return res"""
-
-test3 = """if x == 5:
-        print("x is 5")
-    return res"""
-
-DUMMY_FUNC_SIG = "def func():"
-DUMMY_FUNC_CALL = "func()"
-
-def handle_first_line_indent(func_body: str) -> str:
-    if func_body.startswith("    "):
-        return func_body
-    split = func_body.splitlines()
-    return f"    {split[0]}\n" + "\n".join(split[1:])
-
-def handle_entire_body_indent(func_body: str) -> str:
-    split = func_body.splitlines()
-    res = "\n".join(["    " + line for line in split])
-    return res
-
-def parse_indent(func_body: str) -> str:
-    """
-    3 cases:
-        1. good syntax
-        2. first line not good
-        3. entire body not good
-    """
-    def parse_indent_rec(f_body: str, cur_state: int) -> str:
-        if cur_state > 1:
-            return f_body
-        code = f'{DUMMY_FUNC_SIG}\n{f_body}\n{DUMMY_FUNC_CALL}'
-        try:
-            exec(code)
-            return f_body
-        except (IndentationError, SyntaxError):
-            p_func = handle_first_line_indent if cur_state == 0 else handle_entire_body_indent
-            return parse_indent_rec(p_func(func_body), cur_state + 1)
-        except Exception as e:
-            print(e.args)
-            # print kind of error
-            print(e.__class__.__name__)
-            return f_body
-    return parse_indent_rec(func_body, 0)
-
-# for testing
-if __name__ == "__main__":
-    print(parse_indent(test0))
-    print("\n\n\n")
-    print(parse_indent(test1))
-    print("\n\n\n")
-    print(parse_indent(test2))
-    print("\n\n\n")
-    print(parse_indent(test3))

programming_runs/generators/py_generate.py

@@ -5,6 +5,7 @@ from .generator_utils import generic_generate_func_impl, generic_generate_intern
 from typing import Optional, List, Union
 import ast
 import re
+from parse import parse_code_block, add_code_block
 
 PY_SIMPLE_COMPLETION_INSTRUCTION = "# Write the body of this function only."
 PY_REFLEXION_COMPLETION_INSTRUCTION = "You are a Python writing assistant. You will be given your past function implementation, a series of unit tests, and a hint to change the implementation appropriately. Write your full implementation (restate the function signature).\n\n-----"
@@ -282,7 +283,8 @@ class PyGenerator(Generator):
             simple_chat_instruction=PY_SIMPLE_CHAT_INSTRUCTION,
             reflexion_completion_instruction=PY_REFLEXION_COMPLETION_INSTRUCTION,
             simple_completion_instruction=PY_SIMPLE_COMPLETION_INSTRUCTION,
-            fix_body=fix_turbo_response if strategy == "simple" else py_fix_indentation
+            parse_code_block=lambda x: parse_code_block(x, "python"),
+            add_code_block=lambda x: add_code_block(x, "python"),
         )
 
     def internal_tests(self, func_sig: str, model: ModelBase, max_num_tests: int = 5) -> List[str]:

programming_runs/generators/rs_generate.py

@@ -1,6 +1,7 @@
 from generators.model import ModelBase
 from .generator_types import Generator
 from .generator_utils import generic_generate_func_impl, generic_generate_internal_tests, generic_generate_self_reflection
+from parse import parse_code_block, add_code_block
 
 from typing import List, Optional, Union
 
@@ -176,7 +177,8 @@ class RsGenerator(Generator):
             reflexion_completion_instruction=RS_REFLEXION_COMPLETION_INSTRUCTION,
             simple_completion_instruction=RS_SIMPLE_COMPLETION_INSTRUCTION,
             reflexion_few_shot=RS_REFLEXION_FEW_SHOT_ADD,
-            fix_body=(lambda x: x)
+            parse_code_block=lambda x: parse_code_block(x, "rust"),
+            add_code_block=lambda x: add_code_block(x, "rust"),
         )
 
     def internal_tests(