codingutf8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license testing suite for the pytorch gptneoxjapanese model import unittest from transformers import gptneoxjapaneseconfig istorchavailable from transformers models gptneoxjapanese tokenizationgptneoxjapanese import gptneoxjapanesetokenizer from transformers testingutils import requiretorch slow torchdevice from testconfigurationcommon import configtester from testmodelingcommon import modeltestermixin idstensor randomattentionmask from testpipelinemixin import pipelinetestermixin if istorchavailable import torch from transformers import gptneoxjapaneseforcausallm gptneoxjapanesemodel class gptneoxjapanesemodeltester def init self parent batchsize13 seqlength7 istrainingtrue useinputmasktrue usetokentypeidstrue uselabelstrue vocabsize99 hiddensize32 numhiddenlayers2 numattentionheads4 intermediatemultiplesize4 hiddenactgelu hiddendropout0 0 attentiondropout0 1 weighttyingtrue maxpositionembeddings512 typevocabsize16 typesequencelabelsize2 initializerrange0 02 numlabels3 numchoices4 scopenone self parent parent self batchsize batchsize self seqlength seqlength self istraining istraining self useinputmask useinputmask self usetokentypeids usetokentypeids self uselabels uselabels self vocabsize vocabsize self hiddensize hiddensize self numhiddenlayers numhiddenlayers self numattentionheads numattentionheads self intermediatemultiplesize intermediatemultiplesize self hiddenact hiddenact self hiddendropout hiddendropout self attentiondropout attentiondropout self weighttying weighttying self maxpositionembeddings maxpositionembeddings self typevocabsize typevocabsize self typesequencelabelsize typesequencelabelsize self initializerrange initializerrange self numlabels numlabels self numchoices numchoices self scope scope def prepareconfigandinputsself inputids idstensorself batchsize self seqlength self vocabsize inputmask none if self useinputmask inputmask randomattentionmaskself batchsize self seqlength tokenlabels none if self uselabels tokenlabels idstensorself batchsize self seqlength self numlabels config self getconfig return config inputids inputmask tokenlabels def getconfigself return gptneoxjapaneseconfig vocabsizeself vocabsize hiddensizeself hiddensize numhiddenlayersself numhiddenlayers numattentionheadsself numattentionheads intermediatemultiplesizeself intermediatemultiplesize hiddenactself hiddenact hiddendropoutself hiddendropout attentiondropoutself attentiondropout weighttyingself weighttying maxpositionembeddingsself maxpositionembeddings typevocabsizeself typevocabsize isdecoderfalse initializerrangeself initializerrange def prepareconfigandinputsfordecoderself config inputids inputmask tokenlabels self prepareconfigandinputs config isdecoder true return config inputids inputmask tokenlabels def createandcheckmodelself config inputids inputmask model gptneoxjapanesemodelconfigconfig model totorchdevice model eval modelinputids attentionmaskinputmask result modelinputids self parent assertequalresult lasthiddenstate shape self batchsize self seqlength self hiddensize def createandcheckmodelasdecoderself config inputids inputmask config addcrossattention true model gptneoxjapanesemodelconfig model totorchdevice model eval result modelinputids attentionmaskinputmask self parent assertequalresult lasthiddenstate shape self batchsize self seqlength self hiddensize def createandcheckforcausallmself config inputids inputmask tokenlabels model gptneoxjapaneseforcausallmconfigconfig model totorchdevice model eval result modelinputids attentionmaskinputmask labelstokenlabels self parent assertequalresult logits shape self batchsize self seqlength self vocabsize def createandcheckdecodermodelpastlargeinputsself config inputids inputmask config isdecoder true model gptneoxjapaneseforcausallmconfigconfig model totorchdevice model eval first forward pass outputs modelinputids attentionmaskinputmask usecachetrue pastkeyvalues outputs pastkeyvalues create hypothetical multiple next token and extent to nextinputids nexttokens idstensorself batchsize 3 config vocabsize nextmask idstensorself batchsize 3 vocabsize2 append to next inputids and nextinputids torch catinputids nexttokens dim1 nextattentionmask torch catinputmask nextmask dim1 outputfromnopast modelnextinputids attentionmasknextattentionmask outputhiddenstatestrue outputfromnopast outputfromnopasthiddenstates0 outputfrompast model nexttokens attentionmasknextattentionmask pastkeyvaluespastkeyvalues outputhiddenstatestrue hiddenstates0 select random slice randomsliceidx idstensor1 outputfrompast shape1 item outputfromnopastslice outputfromnopast 3 randomsliceidx detach outputfrompastslice outputfrompast randomsliceidx detach self parent asserttrueoutputfrompastslice shape1 nexttokens shape1 test that outputs are equal for slice self parent asserttruetorch allcloseoutputfrompastslice outputfromnopastslice atol1e3 def prepareconfigandinputsforcommonself configandinputs self prepareconfigandinputs config inputids inputmask tokenlabels configandinputs inputsdict inputids inputids attentionmask inputmask return config inputsdict requiretorch class gptneoxmodeljapanesetestmodeltestermixin pipelinetestermixin unittest testcase allmodelclasses gptneoxjapanesemodel gptneoxjapaneseforcausallm if istorchavailable else allgenerativemodelclasses gptneoxjapaneseforcausallm if istorchavailable else pipelinemodelmapping featureextraction gptneoxjapanesemodel textgeneration gptneoxjapaneseforcausallm if istorchavailable else testpruning false testmissingkeys false testmodelparallel false testheadmasking false def setupself self modeltester gptneoxjapanesemodeltesterself self configtester configtesterself configclassgptneoxjapaneseconfig hiddensize37 def testconfigself self configtester runcommontests def testmodelself config inputids inputmask tokenlabels self modeltester prepareconfigandinputs self modeltester createandcheckmodelconfig inputids inputmask def testmodelasdecoderself config inputids inputmask tokenlabels self modeltester prepareconfigandinputsfordecoder self modeltester createandcheckmodelasdecoderconfig inputids inputmask def testmodelasdecoderwithdefaultinputmaskself this regression test was failing with pytorch 1 3 config inputids inputmask tokenlabels self modeltester prepareconfigandinputsfordecoder inputmask none self modeltester createandcheckmodelasdecoderconfig inputids inputmask def testdecodermodelpastlargeinputsself config inputids inputmask tokenlabels self modeltester prepareconfigandinputs self modeltester createandcheckdecodermodelpastlargeinputsconfig inputids inputmask def testmodelforcausallmself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckforcausallmconfigandinputs slow def testgenerationself modelid abejagptneoxjapanese2 7b prompts 100 fmt skip expectedoutputs 100 tokenizer gptneoxjapanesetokenizer frompretrainedmodelid model gptneoxjapaneseforcausallm frompretrainedmodelid predictedoutputs for prompt in prompts inputids tokenizerprompt returntensorspt inputids generatedids model generateinputids maxlength50 generatedstring tokenizer batchdecodegeneratedids skipspecialtokenstrue predictedoutputs generatedstring self assertlistequalpredictedoutputs expectedoutputs coding utf 8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license testing suite for the pytorch gptneoxjapanese model first forward pass create hypothetical multiple next token and extent to next_input_ids append to next input_ids and select random slice test that outputs are equal for slice this regression test was failing with pytorch 1 3 fmt skip
import unittest

from transformers import GPTNeoXJapaneseConfig, is_torch_available
from transformers.models.gpt_neox_japanese.tokenization_gpt_neox_japanese import GPTNeoXJapaneseTokenizer
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import GPTNeoXJapaneseForCausalLM, GPTNeoXJapaneseModel


class GPTNeoXJapaneseModelTester:
    def __init__(
        self,
        parent,
        batch_size=13,
        seq_length=7,
        is_training=True,
        use_input_mask=True,
        use_token_type_ids=True,
        use_labels=True,
        vocab_size=99,
        hidden_size=32,
        num_hidden_layers=2,
        num_attention_heads=4,
        intermediate_multiple_size=4,
        hidden_act="gelu",
        hidden_dropout=0.0,
        attention_dropout=0.1,
        weight_tying=True,
        max_position_embeddings=512,
        type_vocab_size=16,
        type_sequence_label_size=2,
        initializer_range=0.02,
        num_labels=3,
        num_choices=4,
        scope=None,
    ):
        self.parent = parent
        self.batch_size = batch_size
        self.seq_length = seq_length
        self.is_training = is_training
        self.use_input_mask = use_input_mask
        self.use_token_type_ids = use_token_type_ids
        self.use_labels = use_labels
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.intermediate_multiple_size = intermediate_multiple_size
        self.hidden_act = hidden_act
        self.hidden_dropout = hidden_dropout
        self.attention_dropout = attention_dropout
        self.weight_tying = weight_tying
        self.max_position_embeddings = max_position_embeddings
        self.type_vocab_size = type_vocab_size
        self.type_sequence_label_size = type_sequence_label_size
        self.initializer_range = initializer_range
        self.num_labels = num_labels
        self.num_choices = num_choices
        self.scope = scope

    def prepare_config_and_inputs(self):
        input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)

        input_mask = None
        if self.use_input_mask:
            input_mask = random_attention_mask([self.batch_size, self.seq_length])

        token_labels = None
        if self.use_labels:
            token_labels = ids_tensor([self.batch_size, self.seq_length], self.num_labels)

        config = self.get_config()

        return config, input_ids, input_mask, token_labels

    def get_config(self):
        return GPTNeoXJapaneseConfig(
            vocab_size=self.vocab_size,
            hidden_size=self.hidden_size,
            num_hidden_layers=self.num_hidden_layers,
            num_attention_heads=self.num_attention_heads,
            intermediate_multiple_size=self.intermediate_multiple_size,
            hidden_act=self.hidden_act,
            hidden_dropout=self.hidden_dropout,
            attention_dropout=self.attention_dropout,
            weight_tying=self.weight_tying,
            max_position_embeddings=self.max_position_embeddings,
            type_vocab_size=self.type_vocab_size,
            is_decoder=False,
            initializer_range=self.initializer_range,
        )

    def prepare_config_and_inputs_for_decoder(self):
        config, input_ids, input_mask, token_labels = self.prepare_config_and_inputs()

        config.is_decoder = True

        return config, input_ids, input_mask, token_labels

    def create_and_check_model(self, config, input_ids, input_mask):
        model = GPTNeoXJapaneseModel(config=config)
        model.to(torch_device)
        model.eval()
        _ = model(input_ids, attention_mask=input_mask)
        result = model(input_ids)
        self.parent.assertEqual(result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size))

    def create_and_check_model_as_decoder(self, config, input_ids, input_mask):
        config.add_cross_attention = True
        model = GPTNeoXJapaneseModel(config)
        model.to(torch_device)
        model.eval()
        result = model(input_ids, attention_mask=input_mask)
        self.parent.assertEqual(result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size))

    def create_and_check_for_causal_lm(self, config, input_ids, input_mask, token_labels):
        model = GPTNeoXJapaneseForCausalLM(config=config)
        model.to(torch_device)
        model.eval()
        result = model(input_ids, attention_mask=input_mask, labels=token_labels)
        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.vocab_size))

    def create_and_check_decoder_model_past_large_inputs(self, config, input_ids, input_mask):
        config.is_decoder = True
        model = GPTNeoXJapaneseForCausalLM(config=config)
        model.to(torch_device)
        model.eval()

        
        outputs = model(input_ids, attention_mask=input_mask, use_cache=True)
        past_key_values = outputs.past_key_values

        
        next_tokens = ids_tensor((self.batch_size, 3), config.vocab_size)
        next_mask = ids_tensor((self.batch_size, 3), vocab_size=2)

        
        next_input_ids = torch.cat([input_ids, next_tokens], dim=-1)
        next_attention_mask = torch.cat([input_mask, next_mask], dim=-1)

        output_from_no_past = model(next_input_ids, attention_mask=next_attention_mask, output_hidden_states=True)
        output_from_no_past = output_from_no_past["hidden_states"][0]
        output_from_past = model(
            next_tokens,
            attention_mask=next_attention_mask,
            past_key_values=past_key_values,
            output_hidden_states=True,
        )["hidden_states"][0]

        
        random_slice_idx = ids_tensor((1,), output_from_past.shape[-1]).item()
        output_from_no_past_slice = output_from_no_past[:, -3:, random_slice_idx].detach()
        output_from_past_slice = output_from_past[:, :, random_slice_idx].detach()

        self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1])

        
        self.parent.assertTrue(torch.allclose(output_from_past_slice, output_from_no_past_slice, atol=1e-3))

    def prepare_config_and_inputs_for_common(self):
        config_and_inputs = self.prepare_config_and_inputs()
        config, input_ids, input_mask, token_labels = config_and_inputs
        inputs_dict = {"input_ids": input_ids, "attention_mask": input_mask}
        return config, inputs_dict


@require_torch
class GPTNeoXModelJapaneseTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
    all_model_classes = (GPTNeoXJapaneseModel, GPTNeoXJapaneseForCausalLM) if is_torch_available() else ()
    all_generative_model_classes = (GPTNeoXJapaneseForCausalLM,) if is_torch_available() else ()
    pipeline_model_mapping = (
        {"feature-extraction": GPTNeoXJapaneseModel, "text-generation": GPTNeoXJapaneseForCausalLM}
        if is_torch_available()
        else {}
    )
    test_pruning = False
    test_missing_keys = False
    test_model_parallel = False
    test_head_masking = False

    def setUp(self):
        self.model_tester = GPTNeoXJapaneseModelTester(self)
        self.config_tester = ConfigTester(self, config_class=GPTNeoXJapaneseConfig, hidden_size=37)

    def test_config(self):
        self.config_tester.run_common_tests()

    def test_model(self):
        config, input_ids, input_mask, token_labels = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(config, input_ids, input_mask)

    def test_model_as_decoder(self):
        config, input_ids, input_mask, token_labels = self.model_tester.prepare_config_and_inputs_for_decoder()
        self.model_tester.create_and_check_model_as_decoder(config, input_ids, input_mask)

    def test_model_as_decoder_with_default_input_mask(self):
        
        config, input_ids, input_mask, token_labels = self.model_tester.prepare_config_and_inputs_for_decoder()

        input_mask = None

        self.model_tester.create_and_check_model_as_decoder(config, input_ids, input_mask)

    def test_decoder_model_past_large_inputs(self):
        config, input_ids, input_mask, token_labels = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_decoder_model_past_large_inputs(config, input_ids, input_mask)

    def test_model_for_causal_lm(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_causal_lm(*config_and_inputs)

    @slow
    def test_generation(self):
        model_id = "abeja/gpt-neox-japanese-2.7b"

        prompts = ["データサイエンティストとは、", "100年後に必要とされる会社は、", "フルリモートの環境で働くために必要なことは、", "国境の長いトンネルを抜けると", "美味しい日本食といえば、"]  

        EXPECTED_OUTPUTS = [
            "データサイエンティストとは、データを分析し、ビジネスに役立つ知見を導き出す専門家のことです。",
            "100年後に必要とされる会社は、「人」が中心の会社です。",
            "フルリモートの環境で働くために必要なことは、「自分の時間をコントロールする」ことです。",
            "国境の長いトンネルを抜けると、そこは雪国だった。",
            "美味しい日本食といえば、やっぱりお寿司ですよね。",
        ]

        tokenizer = GPTNeoXJapaneseTokenizer.from_pretrained(model_id)
        model = GPTNeoXJapaneseForCausalLM.from_pretrained(model_id)

        predicted_outputs = []
        for prompt in prompts:
            input_ids = tokenizer(prompt, return_tensors="pt").input_ids
            generated_ids = model.generate(input_ids, max_length=50)
            generated_string = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
            predicted_outputs += generated_string
        self.assertListEqual(predicted_outputs, EXPECTED_OUTPUTS)