prepend and parse prefix cli arg correctly when doing FIM

Author: benlipkin

lm_eval/utils.py

@@ -46,7 +46,9 @@ def __iter__(self):
             elif isinstance(prompt_contents, dict):
                 assert set(prompt_contents.keys()) == {"prefix", "suffix"}
                 infill.append(True)
-                prompt = self.prefix + self._make_infill_prompt(**prompt_contents)
+                prompt = self._make_infill_prompt(
+                    **prompt_contents, preprefix=self.prefix
+                )
             else:
                 raise ValueError(f"Unsupported prompt format: {type(prompt_contents)}")
             prompts.append(prompt)
@@ -83,18 +85,18 @@ def __iter__(self):
                     "input_len": outputs.attention_mask[sample].sum(),
                 }
 
-    def _make_infill_prompt(self, prefix, suffix):
+    def _make_infill_prompt(self, prefix, suffix, preprefix=""):
         """Make a prompt for infilling.
         Currently supported only for official InCoder and SantaCoder implementations.
         """
         model_id = self.tokenizer.name_or_path
         if model_id in ["facebook/incoder-1B", "facebook/incoder-6B"]:
             self.tokenizer.add_special_tokens({"pad_token": "<pad>"})
-            return f"{prefix}<|mask:0|>{suffix}<|mask:0|>"
+            return f"{preprefix}{prefix}<|mask:0|>{suffix}<|mask:0|>"
         elif model_id in ["bigcode/santacoder"]:
-            return f"<fim-prefix>{prefix}<fim-suffix>{suffix}<fim-middle>"
-        elif model_id in ["bigcode/large-model"]:
-            return f"<fim_prefix>{prefix}<fim_suffix>{suffix}<fim_middle>"
+            return f"<fim-prefix>{preprefix}{prefix}<fim-suffix>{suffix}<fim-middle>"
+        elif model_id in ["bigcode/large-model", "bigcode/temp-model"]:
+            return f"<fim_prefix>{preprefix}{prefix}<fim_suffix>{suffix}<fim_middle>"
         else:
             raise ValueError(f"Infilling not yet supported for: {model_id}")
 
@@ -160,7 +162,7 @@ def parse_infill(code, tokenizer):
             prefix, rest = code.split("<fim-suffix>", 1)
             suffix, infill = rest.split("<fim-middle>", 1)
             infill = infill.split("<|endoftext|>")[0]
-        elif model_id in ["bigcode/large-model"]:
+        elif model_id in ["bigcode/large-model", "bigcode/temp-model"]:
             prefix, rest = code.split("<fim_suffix>", 1)
             suffix, infill = rest.split("<fim_middle>", 1)
             infill = infill.split("<|endoftext|>")[0]
@@ -193,9 +195,16 @@ def parse_infill(code, tokenizer):
                     s, skip_special_tokens=True, clean_up_tokenization_spaces=True
                 )
             if postprocess:
-                code_gens[sample].append(
-                    task.postprocess_generation(gen_code[len(prefix) :], int(sample))
-                )
+                if INFILL_MODE:
+                    code_gens[sample].append(
+                        task.postprocess_generation(gen_code, int(sample))
+                    )
+                else:
+                    code_gens[sample].append(
+                        task.postprocess_generation(
+                            gen_code[len(prefix) :], int(sample)
+                        )
+                    )
             else:
                 warnings.warn(
                     "model output is not postprocessed, this might lower evaluation scores"