handle hpc-coder-v2 in generation scripts (#28)

Author: Dando18

generate/generate.py

@@ -130,11 +130,13 @@
         cur_prompt = prompt.copy()
         cur_prompt.update({"temperature": args.temperature, "top_p": args.top_p, "do_sample": args.do_sample, "max_new_tokens": args.max_new_tokens, "prompted": args.prompted})
         cur_prompt["outputs"] = []
+        cur_prompt["raw_outputs"] = []
         prompt_str = cur_prompt["prompt"]
 
     total_tokens += len(generator.tokenizer.encode(output[0]["generated_text"]))
     cleaned_output = inference_config.clean_output(output[0]["generated_text"], prompt_str)
     cur_prompt["outputs"].append(cleaned_output)
+    cur_prompt["raw_outputs"].append(output[0]["generated_text"])
 
     if idx % args.num_samples_per_prompt == args.num_samples_per_prompt - 1:
         responses.append(cur_prompt)

generate/utils.py

@@ -66,6 +66,52 @@ def find_matching_brace_index(code: str, open_brace_index: int) -> int:
     raise ValueError("Unmatched opening brace")
 
 
+def clean_instruct_output(output: str, prompt: str, response_tag: str) -> str:
+    """ Clean LLM output to find code solution. The output should be in a ```c++ ``` code block. If there are
+        multiple, then it tries to find the block with the function definition (as contained in the prompt).
+        The code block itself may include the function definition and body OR just the body. This will try
+        to parse both.
+    """
+    # 0. replace up to the end of the first instance of prompt
+    prompt_loc = output.find(response_tag)
+    if prompt_loc == -1:
+        raise ValueError(f"Prompt not found in output: {prompt}")
+    output = output[prompt_loc + len(response_tag):].strip()
+
+    # 1. Find all code blocks enclosed in triple backticks with "c++" language tag
+    code_blocks = re.findall(r"```\n(.*?)\n```", output, flags=re.DOTALL)
+    code_blocks = [block.removeprefix("```").removeprefix("cpp").removeprefix('c++').removesuffix('```') for block in code_blocks]
+
+    # 2. Prioritize code blocks containing the function definition from the prompt
+    sub_prompt = prompt.rstrip().removesuffix(response_tag).rstrip().removesuffix("```").split("```")[-1]
+    function_name = get_function_name(sub_prompt, "cuda" if "__global__" in sub_prompt else "serial")
+    prioritized_blocks = [block for block in code_blocks if function_name in block]
+
+    # 3. Choose the first block if multiple match, or any block if none match
+    if len(code_blocks) > 0:
+        selected_block = prioritized_blocks[0] if prioritized_blocks else code_blocks[0]
+    else:
+        if '```' in output: # starts with ```c++ but it didn't finish
+            code_idx = output.find('```')
+            selected_block = output[code_idx:].removeprefix('```')
+        else:
+            selected_block = output
+
+    # 4. Handle cases where the block contains only the function body
+    if function_name not in selected_block:
+        return selected_block
+    else:
+        function_start_index = selected_block.index(function_name)
+        open_brace_index = selected_block.find("{", function_start_index)
+        try:
+            close_brace_index = find_matching_brace_index(selected_block, open_brace_index)
+        except ValueError:
+            close_brace_index = len(selected_block)
+
+        function_body = selected_block[open_brace_index + 1 : close_brace_index]
+        return function_body + "}"
+
+
 class InferenceConfig(ABC):
 
     def __init__(self, prompted : bool = False):
@@ -290,50 +336,70 @@ def format_prompt(self, prompt : str) -> str:
         return prompt.strip()
 
     def clean_output(self, output: str, prompt: str) -> str:
-        """ Clean LLM output to find code solution. The output should be in a ```c++ ``` code block. If there are
-            multiple, then it tries to find the block with the function definition (as contained in the prompt).
-            The code block itself may include the function definition and body OR just the body. This will try
-            to parse both.
-        """
-        # 0. replace up to the end of the first instance of prompt
-        prompt_loc = output.find("@@ Response")
-        if prompt_loc == -1:
-            raise ValueError(f"Prompt not found in output: {prompt}")
-        output = output[prompt_loc + len("@@ Response"):].strip()
-
-        # 1. Find all code blocks enclosed in triple backticks with "c++" language tag
-        code_blocks = re.findall(r"```c\+\+\n(.*?)\n```", output, flags=re.DOTALL)
-        code_blocks = [block.lstrip('```c++').rstrip('```') for block in code_blocks]
-
-        # 2. Prioritize code blocks containing the function definition from the prompt
-        sub_prompt = prompt.rstrip().removesuffix("@@ Response").rstrip().removesuffix("```").split("```")[-1]
-        function_name = get_function_name(sub_prompt, "cuda" if "__global__" in sub_prompt else "serial")
-        prioritized_blocks = [block for block in code_blocks if function_name in block]
-
-        # 3. Choose the first block if multiple match, or any block if none match
-        if len(code_blocks) > 0:
-            selected_block = prioritized_blocks[0] if prioritized_blocks else code_blocks[0]
-        else:
-            if '```c++' in output: # starts with ```c++ but it didn't finish
-                code_idx = output.find('```c++')
-                selected_block = output[code_idx:].removeprefix('```c++')
-            else:
-                selected_block = output
+        return clean_instruct_output(output, prompt, "@@ Response")
 
-        # 4. Handle cases where the block contains only the function body
-        if function_name not in selected_block:
-            return selected_block
-        else:
-            function_start_index = selected_block.index(function_name)
-            open_brace_index = selected_block.find("{", function_start_index)
-            try:
-                close_brace_index = find_matching_brace_index(selected_block, open_brace_index)
-            except ValueError:
-                close_brace_index = len(selected_block)
 
-            function_body = selected_block[open_brace_index + 1 : close_brace_index]
-            return function_body + "}"
+class DeepSeekBaseConfig(InferenceConfig):
+
+    def __init__(self, prompted : bool = False):
+        super().__init__(prompted=prompted)
+
+    def get_dtype(self):
+        return torch.bfloat16
+
+    def init_padding(self, tokenizer):
+        tokenizer.pad_token_id = tokenizer.eos_token_id  # for batching
+        tokenizer.padding_side = "left"   # for decoder-only models
 
+    def get_pad_token_id(self, tokenizer) -> int:
+        return tokenizer.pad_token_id
+
+    def get_eos_token_id(self, tokenizer) -> int:
+        return tokenizer.eos_token_id
+    
+    def trust_remote_code(self) -> bool:
+        return False
+
+    def format_prompt(self, prompt : str) -> str:
+        if self.prompted:
+            return f"// filename: solutions/solution_1.cpp\n// here is the correct implementation of the coding exercise\n\n{prompt}"
+        return prompt.strip()
+
+    def clean_output(self, output: str, prompt: str) -> str:
+        return clean_output(output, prompt)
+
+
+class InstructConfig(InferenceConfig):
+
+    def __init__(self, prompted : bool = False, instruction_tag : str = "### Instruction", response_tag : str = "### Response"):
+        super().__init__(prompted=prompted)
+        self.instruction_tag = instruction_tag
+        self.response_tag = response_tag
+
+    def get_dtype(self):
+        return torch.bfloat16
+
+    def init_padding(self, tokenizer):
+        tokenizer.pad_token_id = tokenizer.eos_token_id  # for batching
+        tokenizer.padding_side = "left"   # for decoder-only models
+
+    def get_pad_token_id(self, tokenizer) -> int:
+        return tokenizer.pad_token_id
+
+    def get_eos_token_id(self, tokenizer) -> int:
+        return tokenizer.eos_token_id
+    
+    def trust_remote_code(self) -> bool:
+        return False
+
+    def format_prompt(self, prompt : str) -> str:
+        function_name = get_function_name(prompt, "cuda" if "__global__" in prompt else "serial")
+        prompt = f"Complete the following c++ function.\n```c++{prompt.strip()}```\nWrite only the function {function_name} and no other code. Enclose your solution in ```c++ and ```."
+        prompt = f"{self.instruction_tag}\n{prompt}\n{self.response_tag}\n"
+        return prompt.strip()
+
+    def clean_output(self, output: str, prompt: str) -> str:
+        return clean_instruct_output(output, prompt, self.response_tag)
 
 def get_inference_config(model_name : str, **kwargs) -> InferenceConfig:
     if model_name == "bigcode/starcoderbase":
@@ -350,6 +416,12 @@ def get_inference_config(model_name : str, **kwargs) -> InferenceConfig:
         return ReplitConfig(**kwargs)
     elif model_name.startswith('ise-uiuc/Magicoder'):
         return MagicoderConfig(**kwargs)
+    elif model_name in ['deepseek-ai/deepseek-coder-6.7b-base', 'deepseek-ai/deepseek-coder-7b-base-v1.5']:
+        return DeepSeekBaseConfig(**kwargs)
+    elif model_name.startswith('hpcgroup/hpc-coder-v2'):
+        return InstructConfig(instruction_tag='Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:', response_tag='### Response:', **kwargs)
+    elif model_name.startswith('hpcgroup/rlpf'):
+        return InstructConfig(instruction_tag='### Instruction', response_tag='### Response', **kwargs)
     else:
         raise ValueError(f"Unknown model name: {model_name}")