Merge pull request #66 from benlipkin/BigCodeFIM

added FIM tokens for bigcode/large-model

Author: loubnabnl

lm_eval/utils.py

@@ -93,6 +93,8 @@ def _make_infill_prompt(self, prefix, suffix):
             return f"{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>"
         else:
             raise ValueError(f"Infilling not yet supported for: {model_id}")
 
@@ -158,6 +160,10 @@ 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"]:
+            prefix, rest = code.split("<fim_suffix>", 1)
+            suffix, infill = rest.split("<fim_middle>", 1)
+            infill = infill.split("<|endoftext|>")[0]
         else:
             raise ValueError(f"Infilling not yet supported for: {model_id}")
         for k, v in tokenizer.special_tokens_map.items():