openbmb/cpm-bee-10b

1# coding=utf-82# Copyright 2022 The HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15""" Testing suite for the PyTorch CpmBee model. """16 17 18import unittest19 20from transformers.testing_utils import is_torch_available, require_torch, tooslow21 22from ...generation.test_utils import torch_device23from ...test_configuration_common import ConfigTester24from ...test_modeling_common import ModelTesterMixin, ids_tensor25from ...test_pipeline_mixin import PipelineTesterMixin26 27 28if is_torch_available():29    import torch30 31    from transformers import (32        CpmBeeConfig,33        CpmBeeForCausalLM,34        CpmBeeModel,35        CpmBeeTokenizer,36    )37 38 39@require_torch40class CpmBeeModelTester:41    def __init__(42        self,43        parent,44        batch_size=2,45        seq_length=8,46        is_training=True,47        use_token_type_ids=False,48        use_input_mask=False,49        use_labels=False,50        use_mc_token_ids=False,51        vocab_size=99,52        hidden_size=32,53        num_hidden_layers=3,54        num_attention_heads=4,55        intermediate_size=37,56        num_buckets=32,57        max_distance=128,58        position_bias_num_segment_buckets=32,59        init_std=1.0,60        return_dict=True,61    ):62        self.parent = parent63        self.batch_size = batch_size64        self.seq_length = seq_length65        self.is_training = is_training66        self.use_token_type_ids = use_token_type_ids67        self.use_input_mask = use_input_mask68        self.use_labels = use_labels69        self.use_mc_token_ids = use_mc_token_ids70        self.vocab_size = vocab_size71        self.hidden_size = hidden_size72        self.num_hidden_layers = num_hidden_layers73        self.num_attention_heads = num_attention_heads74        self.intermediate_size = intermediate_size75        self.num_buckets = num_buckets76        self.max_distance = max_distance77        self.position_bias_num_segment_buckets = position_bias_num_segment_buckets78        self.init_std = init_std79        self.return_dict = return_dict80 81    def prepare_config_and_inputs(self):82        input_ids = {}83        input_ids["input_ids"] = ids_tensor([self.batch_size, self.seq_length], self.vocab_size).type(torch.int32)84        input_ids["use_cache"] = False85 86        config = self.get_config()87 88        return (config, input_ids)89 90    def get_config(self):91        return CpmBeeConfig(92            vocab_size=self.vocab_size,93            hidden_size=self.hidden_size,94            num_hidden_layers=self.num_hidden_layers,95            num_attention_heads=self.num_attention_heads,96            dim_ff=self.intermediate_size,97            position_bias_num_buckets=self.num_buckets,98            position_bias_max_distance=self.max_distance,99            position_bias_num_segment_buckets=self.position_bias_num_segment_buckets,100            use_cache=True,101            init_std=self.init_std,102            return_dict=self.return_dict,103        )104 105    def create_and_check_cpmbee_model(self, config, input_ids, *args):106        model = CpmBeeModel(config=config)107        model.to(torch_device)108        model.eval()109 110        hidden_states = model(**input_ids).last_hidden_state111 112        self.parent.assertEqual(hidden_states.shape, (self.batch_size, self.seq_length, config.hidden_size))113 114    def create_and_check_lm_head_model(self, config, input_ids, *args):115        model = CpmBeeForCausalLM(config)116        model.to(torch_device)117        input_ids["input_ids"] = input_ids["input_ids"].to(torch_device)118        model.eval()119 120        model_output = model(**input_ids)121        self.parent.assertEqual(122            model_output.logits.shape,123            (self.batch_size, self.seq_length, config.vocab_size),124        )125 126    def prepare_config_and_inputs_for_common(self):127        config, inputs_dict = self.prepare_config_and_inputs()128        return config, inputs_dict129 130 131@require_torch132class CpmBeeModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):133    all_model_classes = (CpmBeeModel, CpmBeeForCausalLM) if is_torch_available() else ()134    pipeline_model_mapping = (135        {"feature-extraction": CpmBeeModel, "text-generation": CpmBeeForCausalLM} if is_torch_available() else {}136    )137 138    test_pruning = False139    test_missing_keys = False140    test_mismatched_shapes = False141    test_head_masking = False142    test_resize_embeddings = False143 144    def setUp(self):145        self.model_tester = CpmBeeModelTester(self)146        self.config_tester = ConfigTester(self, config_class=CpmBeeConfig)147 148    def test_config(self):149        self.config_tester.create_and_test_config_common_properties()150        self.config_tester.create_and_test_config_to_json_string()151        self.config_tester.create_and_test_config_to_json_file()152        self.config_tester.create_and_test_config_from_and_save_pretrained()153        self.config_tester.check_config_can_be_init_without_params()154        self.config_tester.check_config_arguments_init()155 156    def test_inputs_embeds(self):157        unittest.skip("CPMBee doesn't support input_embeds.")(self.test_inputs_embeds)158 159    def test_retain_grad_hidden_states_attentions(self):160        unittest.skip(161            "CPMBee doesn't support retain grad in hidden_states or attentions, because prompt management will peel off the output.hidden_states from graph.\162                 So is attentions. We strongly recommand you use loss to tune model."163        )(self.test_retain_grad_hidden_states_attentions)164 165    def test_cpmbee_model(self):166        config, inputs = self.model_tester.prepare_config_and_inputs()167        self.model_tester.create_and_check_cpmbee_model(config, inputs)168 169    def test_cpmbee_lm_head_model(self):170        config, inputs = self.model_tester.prepare_config_and_inputs()171        self.model_tester.create_and_check_lm_head_model(config, inputs)172 173 174@require_torch175class CpmBeeForCausalLMlIntegrationTest(unittest.TestCase):176    @tooslow177    def test_simple_generation(self):178        texts = {"input": "今天天气不错，", "<ans>": ""}179        model = CpmBeeForCausalLM.from_pretrained("openbmb/cpm-bee-10b")180        tokenizer = CpmBeeTokenizer.from_pretrained("openbmb/cpm-bee-10b")181        output_texts = model.generate(texts, tokenizer)182        expected_output = {"input": "今天天气不错，", "<ans>": "适合睡觉。"}183        self.assertEqual(expected_output["<ans>"], output_texts["<ans>"])184