diff --git a/programming_runs/generators/generator_utils.py b/programming_runs/generators/generator_utils.py index 1f4798a..dd05736 100644 --- a/programming_runs/generators/generator_utils.py +++ b/programming_runs/generators/generator_utils.py @@ -5,7 +5,6 @@ from typing import Union, List, Optional, Callable # openai.api_key = os.getenv("OPENAI_API_KEY") -USE_PYTHON_CODEBLOCK_INSTRUCTION = "Use a Python code block to write your response. For example:\n```python\nprint('Hello world!')\n```" def generic_generate_func_impl( @@ -22,6 +21,7 @@ def generic_generate_func_impl( simple_chat_instruction: str, reflexion_completion_instruction: str, simple_completion_instruction: str, + code_block_instruction: str, parse_code_block: Callable[[str], str], add_code_block: Callable[[str], str] ) -> Union[str, List[str]]: @@ -35,7 +35,7 @@ def generic_generate_func_impl( if model.is_chat: if strategy == "reflexion": message = f"{reflexion_few_shot}\n[previous impl]:\n{add_code_block(prev_func_impl)}\n\n[unit test results from previous impl]:\n{feedback}\n\n[reflection on previous impl]:\n{self_reflection}\n\n[improved impl]:\n{func_sig}" - prompt = f"{reflexion_chat_instruction}\n{USE_PYTHON_CODEBLOCK_INSTRUCTION}" + prompt = f"{reflexion_chat_instruction}\n{code_block_instruction}" # func_bodies is a really bad name, as it can also be just 1 string print_messages(prompt, message) messages = [ @@ -66,12 +66,12 @@ def generic_generate_func_impl( ] func_bodies = model.generate_chat(messages=messages, num_comps=num_comps, temperature=temperature) else: - system_prompt = f"{simple_chat_instruction}\n{USE_PYTHON_CODEBLOCK_INSTRUCTION}" + system_prompt = f"{simple_chat_instruction}\n{code_block_instruction}" print_messages(system_prompt, func_sig) messages = [ Message( role="system", - content=f"{simple_chat_instruction}\n{USE_PYTHON_CODEBLOCK_INSTRUCTION}", + content=f"{simple_chat_instruction}\n{code_block_instruction}", ), Message( role="user", @@ -81,11 +81,11 @@ def generic_generate_func_impl( func_bodies = model.generate_chat(messages=messages, num_comps=num_comps, temperature=temperature) else: if strategy == "reflexion": - prompt = f"{reflexion_completion_instruction}\n{add_code_block(prev_func_impl)}\n\nunit tests:\n{feedback}\n\nhint:\n{self_reflection}\n\n# improved implementation\n{func_sig}\n{USE_PYTHON_CODEBLOCK_INSTRUCTION}" + prompt = f"{reflexion_completion_instruction}\n{add_code_block(prev_func_impl)}\n\nunit tests:\n{feedback}\n\nhint:\n{self_reflection}\n\n# improved implementation\n{func_sig}\n{code_block_instruction}" func_bodies = model.generate( prompt, num_comps=num_comps, temperature=temperature) else: - prompt = f"{simple_completion_instruction}\n{func_sig}\n{USE_PYTHON_CODEBLOCK_INSTRUCTION}" + prompt = f"{simple_completion_instruction}\n{func_sig}\n{code_block_instruction}" func_bodies = model.generate( prompt, num_comps=num_comps, temperature=temperature) diff --git a/programming_runs/generators/model.py b/programming_runs/generators/model.py index cc94d05..02f173e 100644 --- a/programming_runs/generators/model.py +++ b/programming_runs/generators/model.py @@ -125,7 +125,7 @@ class StarChat(ModelBase): prompt += f"<|{message.role}|>\n{message.content}<|end|>\n" if i != len(messages) - 1: prompt += "\n<|assistant|>" - + outputs = self.pipe( prompt, max_new_tokens=max_tokens, diff --git a/programming_runs/generators/py_generate.py b/programming_runs/generators/py_generate.py index 4701e03..48c354c 100644 --- a/programming_runs/generators/py_generate.py +++ b/programming_runs/generators/py_generate.py @@ -10,6 +10,7 @@ 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-----" PY_SELF_REFLECTION_COMPLETION_INSTRUCTION = "You are a Python writing assistant. You will be given a function implementation and a series of unit tests. Your goal is to write a few sentences to explain why your implementation is wrong as indicated by the tests. You will need this as a hint when you try again later. Only provide the few sentence description in your answer, not the implementation.\n\n-----" +USE_PYTHON_CODEBLOCK_INSTRUCTION = "Use a Python code block to write your response. For example:\n```python\nprint('Hello world!')\n```" PY_SIMPLE_CHAT_INSTRUCTION = "You are an AI that only responds with python code, NOT ENGLISH. You will be given a function signature and its docstring by the user. Write your full implementation (restate the function signature)." PY_SIMPLE_CHAT_INSTRUCTION_V2 = "You are an AI that only responds with only python code. You will be given a function signature and its docstring by the user. Write your full implementation (restate the function signature)." @@ -284,6 +285,7 @@ class PyGenerator(Generator): simple_chat_instruction=PY_SIMPLE_CHAT_INSTRUCTION, reflexion_completion_instruction=PY_REFLEXION_COMPLETION_INSTRUCTION, simple_completion_instruction=PY_SIMPLE_COMPLETION_INSTRUCTION, + code_block_instruction=USE_PYTHON_CODEBLOCK_INSTRUCTION, parse_code_block=lambda x: parse_code_block(x, "python"), add_code_block=lambda x: add_code_block(x, "python"), ) diff --git a/programming_runs/generators/rs_generate.py b/programming_runs/generators/rs_generate.py index 2da41b7..14ccd7a 100644 --- a/programming_runs/generators/rs_generate.py +++ b/programming_runs/generators/rs_generate.py @@ -8,6 +8,7 @@ from typing import List, Optional, Union RS_SIMPLE_COMPLETION_INSTRUCTION = "// Write the body of this function only." RS_REFLEXION_COMPLETION_INSTRUCTION = "You are a Rust 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-----" RS_SELF_REFLECTION_COMPLETION_INSTRUCTION = "You are a Rust writing assistant. You will be given a function implementation and a series of unit tests. Your goal is to write a few sentences to explain why your implementation is wrong as indicated by the tests. You will need this as a hint when you try again later. Only provide the few sentence description in your answer, not the implementation.\n\n-----" +USE_RUST_CODEBLOCK_INSTRUCTION = "Use a Rust code block to write your response. For example:\n```rust\nfn main() {\n println!(\"Hello\");\n}\n```" RS_SIMPLE_CHAT_INSTRUCTION = "You are an AI that only responds with Rust code, NOT ENGLISH. You will be given a function signature and its docstring by the user. Write your full implementation (restate the function signature)." RS_REFLEXION_CHAT_INSTRUCTION = "You are an AI Rust 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)."