fix code block instruction always python

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)
 

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,

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"),
         )

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)."