Support chat model human-eval test.

Author: gushiqiao

configs/quantization/methods/RTN/rtn_w_a_kv_human_eval.yml

@@ -14,6 +14,7 @@ eval:
     bs: 1
     format_tabs: True
     inference_per_block: False
+    # add_chat_temp: True
 quant:
     method: RTN
     weight:

llmc/eval/eval_base.py

@@ -13,29 +13,24 @@ class BaseEval:
     def __init__(self, tokenizer, config):
         self.tokenizer = tokenizer
         # eval_cfg
-        eval_cfg = config.eval
+        self.eval_cfg = config.eval
         self.model_type = config.model.type
-        logger.info(f'eval_cfg : {eval_cfg}')
-        self.dataset = eval_cfg['name']
+        logger.info(f'eval_cfg : {self.eval_cfg}')
+        self.dataset = self.eval_cfg['name']
         assert self.dataset in [
             'wikitext2',
             'c4',
             'ptb',
             'custom',
             'human_eval'
-        ], 'Ppl eval only support wikitext2, c4, ptb, human_eval dataset now.'
-        self.seq_len = eval_cfg.get('seq_len', None)
-        self.bs = eval_cfg['bs']
-        self.path = eval_cfg.get('path', None)
-        self.download = eval_cfg.get('download', False)
-        self.load_from_txt = eval_cfg.get('load_from_txt', False)
-        self.inference_per_block = eval_cfg.get('inference_per_block', False)
+        ], 'Eval only support wikitext2, c4, ptb, custom, human_eval dataset now.'
+        self.seq_len = self.eval_cfg.get('seq_len', None)
+        self.bs = self.eval_cfg['bs']
+        self.path = self.eval_cfg.get('path', None)
+        self.download = self.eval_cfg.get('download', False)
+        self.load_from_txt = self.eval_cfg.get('load_from_txt', False)
+        self.inference_per_block = self.eval_cfg.get('inference_per_block', False)
         self.testenc = self.build_data()
-        self.res_path = eval_cfg.get('res_path', None)
-        if self.dataset in ['human_eval']:
-            assert self.res_path is not None
-            os.makedirs(self.res_path, exist_ok=True)
-        self.format_tabs = eval_cfg.get('format_tabs', False)
 
     @torch.no_grad()
     def build_data(self):

llmc/eval/eval_code.py

@@ -11,6 +11,15 @@
 
 
 class HumanEval(BaseEval):
+    def __init__(self, tokenizer, config):
+        super().__init__(tokenizer, config)
+        self.res_path = self.eval_cfg.get('res_path', None)
+        assert self.res_path is not None
+        os.makedirs(self.res_path, exist_ok=True)
+        self.format_tabs = self.eval_cfg.get('format_tabs', False)
+        self.instruction = self.eval_cfg.get('instruction',
+                                             'Complete the following Python code:')
+        self.add_chat_temp = self.eval_cfg.get('add_chat_temp', False)
 
     @torch.no_grad()
     def eval_func(self, org_model, model, testenc, seq_len, bs, eval_pos):
@@ -22,6 +31,7 @@ def eval_func(self, org_model, model, testenc, seq_len, bs, eval_pos):
                 prompt = testenc[task_id]['prompt'].replace('    ', '\t')
             else:
                 prompt = testenc[task_id]['prompt']
+            prompt = self.gen_prompt(prompt)
             batch_completions = self.generate_batch_completion(
                 model, prompt, bs
             )
@@ -46,8 +56,24 @@ def eval_func(self, org_model, model, testenc, seq_len, bs, eval_pos):
         res = self.post_process(testenc)
         return res
 
+    def gen_prompt(self, prompt):
+        prompt = self.instruction + '\n' + prompt
+        if self.model_type in ['Starcoder']:
+            prompt = '<fim_prefix>' + prompt + '<fim_suffix><fim_middle>'
+
+        if self.add_chat_temp:
+            chat_prompt = [{'role': 'user', 'content': prompt}]
+            chat_prompt = self.tokenizer.apply_chat_template(
+                chat_prompt,
+                tokenize=False,
+                add_generation_prompt=True
+            )
+            return chat_prompt
+
+        return prompt
+
     @torch.no_grad()
-    def generated_llama(
+    def generated(
         self,
         model,
         inputs,
@@ -56,14 +82,20 @@ def generated_llama(
         top_p=0.95,
         do_sample=True,
     ):
+
+        if hasattr(self.tokenizer, 'pad_token_id'):
+            pad_token_id = self.tokenizer.pad_token_id
+        else:
+            pad_token_id = self.tokenizer.eos_token_id
+
         generated_ids = model.model.generate(
             **inputs,
             max_new_tokens=max_new_tokens,
             temperature=temperature,
             top_p=top_p,
             do_sample=do_sample,
             eos_token_id=self.tokenizer.eos_token_id,
-            pad_token_id=self.tokenizer.eos_token_id,
+            pad_token_id=pad_token_id,
             use_cache=True,
         )
         return generated_ids
@@ -74,11 +106,8 @@ def generate_batch_completion(self, model, prompt, bs):
         inputs = self.tokenizer(input_batch, return_tensors='pt').to(model.model.device)
         input_ids_cutoff = inputs.input_ids.size(dim=1)
 
-        if self.model_type in ['Llama']:
-            generated_ids = self.generated_llama(model, inputs)
-            model.reset_kv()
-        else:
-            raise NotImplementedError('This model is not support yet.')
+        generated_ids = self.generated(model, inputs)
+        model.reset_kv()
 
         batch_completions = self.tokenizer.batch_decode(
             [ids[input_ids_cutoff:] for ids in generated_ids],

llmc/models/internlm2.py

@@ -32,6 +32,9 @@ def get_pre_head_layernorm_layers(self):
     def get_layers_except_blocks(self):
         return [self.tok_embeddings, self.model.model.norm, self.model.output]
 
+    def get_attn_in_block(self, block):
+        return {'attention': block.attention}
+
     def skip_layer_name(self):
         return ['lm_head']