codingutf8 2021 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 opt model import copy import tempfile import unittest import timeoutdecorator noqa from transformers import optconfig istorchavailable from transformers testingutils import requiretorch requiretorchaccelerator requiretorchfp16 slow torchdevice from generation testutils import generationtestermixin from testconfigurationcommon import configtester from testmodelingcommon import modeltestermixin idstensor from testpipelinemixin import pipelinetestermixin if istorchavailable import torch from transformers import gpt2tokenizer optforcausallm optforquestionanswering optforsequenceclassification optmodel def prepareoptinputsdict config inputids decoderinputidsnone attentionmasknone decoderattentionmasknone headmasknone decoderheadmasknone if attentionmask is none attentionmask inputids neconfig padtokenid return inputids inputids attentionmask attentionmask headmask headmask class optmodeltester def init self parent batchsize13 seqlength7 istrainingtrue uselabelsfalse vocabsize99 hiddensize16 numhiddenlayers2 numattentionheads4 intermediatesize4 hiddenactgelu hiddendropoutprob0 1 attentionprobsdropoutprob0 1 maxpositionembeddings20 eostokenid2 padtokenid1 bostokenid0 embeddim16 numlabels3 wordembedprojdim16 typesequencelabelsize2 self parent parent self batchsize batchsize self seqlength seqlength self istraining istraining self uselabels uselabels self vocabsize vocabsize self hiddensize hiddensize self numhiddenlayers numhiddenlayers self numattentionheads numattentionheads self intermediatesize intermediatesize self hiddenact hiddenact self hiddendropoutprob hiddendropoutprob self attentionprobsdropoutprob attentionprobsdropoutprob self maxpositionembeddings maxpositionembeddings self eostokenid eostokenid self padtokenid padtokenid self bostokenid bostokenid self embeddim embeddim self numlabels numlabels self typesequencelabelsize typesequencelabelsize self wordembedprojdim wordembedprojdim self isencoderdecoder false def prepareconfigandinputsself inputids idstensorself batchsize self seqlength self vocabsize clamp 3 inputids 1 self eostokenid eos token decoderinputids idstensorself batchsize self seqlength self vocabsize config self getconfig inputsdict prepareoptinputsdictconfig inputids decoderinputids return config inputsdict def getconfigself return optconfig vocabsizeself vocabsize hiddensizeself hiddensize numhiddenlayersself numhiddenlayers numattentionheadsself numattentionheads ffndimself intermediatesize dropoutself hiddendropoutprob attentiondropoutself attentionprobsdropoutprob maxpositionembeddingsself maxpositionembeddings eostokenidself eostokenid bostokenidself bostokenid padtokenidself padtokenid embeddimself embeddim isencoderdecoderfalse wordembedprojdimself wordembedprojdim def getpipelineconfigself config self getconfig config maxpositionembeddings 100 return config def prepareconfigandinputsforcommonself config inputsdict self prepareconfigandinputs return config inputsdict def createandcheckdecodermodelpastlargeinputsself config inputsdict model optmodelconfigconfig totorchdevice eval inputids inputsdictinputids attentionmask inputsdictattentionmask headmask inputsdictheadmask first forward pass outputs modelinputids attentionmaskattentionmask headmaskheadmask usecachetrue output pastkeyvalues outputs totuple create hypothetical multiple next token and extent to nextinputids nexttokens idstensorself batchsize 3 config vocabsize nextattnmask idstensorself batchsize 3 2 append to next inputids and nextinputids torch catinputids nexttokens dim1 nextattentionmask torch catattentionmask nextattnmask dim1 outputfromnopast modelnextinputids attentionmasknextattentionmasklasthiddenstate outputfrompast modelnexttokens attentionmasknextattentionmask pastkeyvaluespastkeyvalues lasthiddenstate 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 test no attentionmask works outputs modelinputids attentionmaskattentionmask headmaskheadmask usecachetrue pastkeyvalues outputs totuple outputfromnopast modelnextinputidslasthiddenstate outputfrompast modelnexttokens pastkeyvaluespastkeyvalueslasthiddenstate randomsliceidx idstensor1 outputfrompast shape1 item outputfromnopastslice outputfromnopast 3 randomsliceidx detach outputfrompastslice outputfrompast randomsliceidx detach test that outputs are equal for slice self parent asserttruetorch allcloseoutputfrompastslice outputfromnopastslice atol1e3 requiretorch class optmodeltestmodeltestermixin generationtestermixin pipelinetestermixin unittest testcase allmodelclasses optmodel optforcausallm optforsequenceclassification optforquestionanswering if istorchavailable else allgenerativemodelclasses optforcausallm if istorchavailable else pipelinemodelmapping featureextraction optmodel questionanswering optforquestionanswering textclassification optforsequenceclassification textgeneration optforcausallm zeroshot optforsequenceclassification if istorchavailable else isencoderdecoder false fxcompatible true testpruning false testmissingkeys false todo fix the failed tests def ispipelinetesttoskip self pipelinetestcassename configclass modelarchitecture tokenizername processorname if pipelinetestcassename qapipelinetests and tokenizername is not none and not tokenizername endswithfast qapipelinetests fails for a few models when the slower tokenizer are used the slower tokenizers were never used for pipeline tests before the pipeline testing rework todo check and possibly fix the qapipelinetests with slower tokenizer return true return false def setupself self modeltester optmodeltesterself self configtester configtesterself configclassoptconfig def testconfigself self configtester runcommontests def testsaveloadstrictself config inputsdict self modeltester prepareconfigandinputs for modelclass in self allmodelclasses model modelclassconfig with tempfile temporarydirectory as tmpdirname model savepretrainedtmpdirname model2 info modelclass frompretrainedtmpdirname outputloadinginfotrue self assertequalinfomissingkeys def testdecodermodelpastwithlargeinputsself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckdecodermodelpastlargeinputsconfigandinputs def testinputsembedsself config inputsdict self modeltester prepareconfigandinputsforcommon for modelclass in optmodel model modelclassconfig model totorchdevice model eval inputs copy deepcopyself prepareforclassinputsdict modelclass if not self isencoderdecoder inputids inputsinputids del inputsinputids else encoderinputids inputsinputids decoderinputids inputs getdecoderinputids encoderinputids del inputsinputids inputs popdecoderinputids none wte model getinputembeddings if not self isencoderdecoder inputsinputsembeds wteinputids else inputsinputsembeds wteencoderinputids inputsdecoderinputsembeds wtedecoderinputids with torch nograd modelinputs0 requiretorchfp16 def testgeneratefp16self config inputdict self modeltester prepareconfigandinputs inputids inputdictinputids attentionmask inputids ne1 totorchdevice model optforcausallmconfig eval totorchdevice model half model generateinputids attentionmaskattentionmask model generatenumbeams4 dosampletrue earlystoppingfalse numreturnsequences3 def testoptsequenceclassificationmodelself config inputdict self modeltester prepareconfigandinputs config numlabels 3 inputids inputdictinputids attentionmask inputids ne1 totorchdevice sequencelabels idstensorself modeltester batchsize self modeltester typesequencelabelsize model optforsequenceclassificationconfig model totorchdevice model eval result modelinputids attentionmaskattentionmask labelssequencelabels self assertequalresult logits shape self modeltester batchsize self modeltester numlabels def testoptsequenceclassificationmodelformultilabelself config inputdict self modeltester prepareconfigandinputs config numlabels 3 config problemtype multilabelclassification inputids inputdictinputids attentionmask inputids ne1 totorchdevice sequencelabels idstensor self modeltester batchsize config numlabels self modeltester typesequencelabelsize totorch float model optforsequenceclassificationconfig model totorchdevice model eval result modelinputids attentionmaskattentionmask labelssequencelabels self assertequalresult logits shape self modeltester batchsize self modeltester numlabels unittest skipdoes not work on the tiny model as we keep hitting edge cases def testmodelparallelismself super testmodelparallelism def asserttensorsclosea b atol1e12 prefix expected value works for cpu as well as gpu with tf32 disabled verify that prompt without bos token is identical to metaseq addspecialtokensfalse logitsmeta torch loadself pathlogitsmeta use different length sentences to test batching a bug manifested starting at models facebookopt1 3 and larger when running batched generations therefore not using a tiny model but the smallest model the problem was seen with which is opt1 3b please refer to this github thread https github comhuggingfacetransformerspull17437 for more details coding utf 8 2021 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 opt model noqa eos token 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 test no attention_mask works test that outputs are equal for slice todo fix the failed tests qapipelinetests fails for a few models when the slower tokenizer are used the slower tokenizers were never used for pipeline tests before the pipeline testing rework todo check and possibly fix the qapipelinetests with slower tokenizer if tensors have different shapes different values or a and b are not both tensors raise a nice assertion error expected value works for cpu as well as gpu with tf32 disabled verify that prompt without bos token is identical to metaseq add_special_tokens false logits_meta torch load self path_logits_meta use different length sentences to test batching a bug manifested starting at models facebook opt 1 3 and larger when running batched generations therefore not using a tiny model but the smallest model the problem was seen with which is opt 1 3b please refer to this github thread https github com huggingface transformers pull 17437 for more details the first logits could contain nans if it fails
import copy
import tempfile
import unittest

import timeout_decorator  

from transformers import OPTConfig, is_torch_available
from transformers.testing_utils import require_torch, require_torch_accelerator, require_torch_fp16, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import (
        GPT2Tokenizer,
        OPTForCausalLM,
        OPTForQuestionAnswering,
        OPTForSequenceClassification,
        OPTModel,
    )


def prepare_opt_inputs_dict(
    config,
    input_ids,
    decoder_input_ids=None,
    attention_mask=None,
    decoder_attention_mask=None,
    head_mask=None,
    decoder_head_mask=None,
):
    if attention_mask is None:
        attention_mask = input_ids.ne(config.pad_token_id)
    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "head_mask": head_mask,
    }


class OPTModelTester:
    def __init__(
        self,
        parent,
        batch_size=13,
        seq_length=7,
        is_training=True,
        use_labels=False,
        vocab_size=99,
        hidden_size=16,
        num_hidden_layers=2,
        num_attention_heads=4,
        intermediate_size=4,
        hidden_act="gelu",
        hidden_dropout_prob=0.1,
        attention_probs_dropout_prob=0.1,
        max_position_embeddings=20,
        eos_token_id=2,
        pad_token_id=1,
        bos_token_id=0,
        embed_dim=16,
        num_labels=3,
        word_embed_proj_dim=16,
        type_sequence_label_size=2,
    ):
        self.parent = parent
        self.batch_size = batch_size
        self.seq_length = seq_length
        self.is_training = is_training
        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_size = intermediate_size
        self.hidden_act = hidden_act
        self.hidden_dropout_prob = hidden_dropout_prob
        self.attention_probs_dropout_prob = attention_probs_dropout_prob
        self.max_position_embeddings = max_position_embeddings
        self.eos_token_id = eos_token_id
        self.pad_token_id = pad_token_id
        self.bos_token_id = bos_token_id
        self.embed_dim = embed_dim
        self.num_labels = num_labels
        self.type_sequence_label_size = type_sequence_label_size
        self.word_embed_proj_dim = word_embed_proj_dim
        self.is_encoder_decoder = False

    def prepare_config_and_inputs(self):
        input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size).clamp(
            3,
        )
        input_ids[:, -1] = self.eos_token_id  

        decoder_input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)

        config = self.get_config()
        inputs_dict = prepare_opt_inputs_dict(config, input_ids, decoder_input_ids)
        return config, inputs_dict

    def get_config(self):
        return OPTConfig(
            vocab_size=self.vocab_size,
            hidden_size=self.hidden_size,
            num_hidden_layers=self.num_hidden_layers,
            num_attention_heads=self.num_attention_heads,
            ffn_dim=self.intermediate_size,
            dropout=self.hidden_dropout_prob,
            attention_dropout=self.attention_probs_dropout_prob,
            max_position_embeddings=self.max_position_embeddings,
            eos_token_id=self.eos_token_id,
            bos_token_id=self.bos_token_id,
            pad_token_id=self.pad_token_id,
            embed_dim=self.embed_dim,
            is_encoder_decoder=False,
            word_embed_proj_dim=self.word_embed_proj_dim,
        )

    def get_pipeline_config(self):
        config = self.get_config()
        config.max_position_embeddings = 100
        return config

    def prepare_config_and_inputs_for_common(self):
        config, inputs_dict = self.prepare_config_and_inputs()
        return config, inputs_dict

    def create_and_check_decoder_model_past_large_inputs(self, config, inputs_dict):
        model = OPTModel(config=config).to(torch_device).eval()

        input_ids = inputs_dict["input_ids"]
        attention_mask = inputs_dict["attention_mask"]
        head_mask = inputs_dict["head_mask"]

        
        outputs = model(input_ids, attention_mask=attention_mask, head_mask=head_mask, use_cache=True)

        output, past_key_values = outputs.to_tuple()

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

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

        output_from_no_past = model(next_input_ids, attention_mask=next_attention_mask)["last_hidden_state"]
        output_from_past = model(next_tokens, attention_mask=next_attention_mask, past_key_values=past_key_values)[
            "last_hidden_state"
        ]

        
        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))

        
        outputs = model(input_ids, attention_mask=attention_mask, head_mask=head_mask, use_cache=True)
        _, past_key_values = outputs.to_tuple()
        output_from_no_past = model(next_input_ids)["last_hidden_state"]

        output_from_past = model(next_tokens, past_key_values=past_key_values)["last_hidden_state"]

        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(torch.allclose(output_from_past_slice, output_from_no_past_slice, atol=1e-3))


@require_torch
class OPTModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, unittest.TestCase):
    all_model_classes = (
        (OPTModel, OPTForCausalLM, OPTForSequenceClassification, OPTForQuestionAnswering)
        if is_torch_available()
        else ()
    )
    all_generative_model_classes = (OPTForCausalLM,) if is_torch_available() else ()
    pipeline_model_mapping = (
        {
            "feature-extraction": OPTModel,
            "question-answering": OPTForQuestionAnswering,
            "text-classification": OPTForSequenceClassification,
            "text-generation": OPTForCausalLM,
            "zero-shot": OPTForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    is_encoder_decoder = False
    fx_compatible = True
    test_pruning = False
    test_missing_keys = False

    
    def is_pipeline_test_to_skip(
        self, pipeline_test_casse_name, config_class, model_architecture, tokenizer_name, processor_name
    ):
        if (
            pipeline_test_casse_name == "QAPipelineTests"
            and tokenizer_name is not None
            and not tokenizer_name.endswith("Fast")
        ):
            
            
            
            return True

        return False

    def setUp(self):
        self.model_tester = OPTModelTester(self)
        self.config_tester = ConfigTester(self, config_class=OPTConfig)

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

    def test_save_load_strict(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs()
        for model_class in self.all_model_classes:
            model = model_class(config)

            with tempfile.TemporaryDirectory() as tmpdirname:
                model.save_pretrained(tmpdirname)
                model2, info = model_class.from_pretrained(tmpdirname, output_loading_info=True)
            self.assertEqual(info["missing_keys"], [])

    def test_decoder_model_past_with_large_inputs(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_decoder_model_past_large_inputs(*config_and_inputs)

    def test_inputs_embeds(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in (OPTModel,):
            model = model_class(config)
            model.to(torch_device)
            model.eval()

            inputs = copy.deepcopy(self._prepare_for_class(inputs_dict, model_class))

            if not self.is_encoder_decoder:
                input_ids = inputs["input_ids"]
                del inputs["input_ids"]
            else:
                encoder_input_ids = inputs["input_ids"]
                decoder_input_ids = inputs.get("decoder_input_ids", encoder_input_ids)
                del inputs["input_ids"]
                inputs.pop("decoder_input_ids", None)

            wte = model.get_input_embeddings()
            if not self.is_encoder_decoder:
                inputs["inputs_embeds"] = wte(input_ids)
            else:
                inputs["inputs_embeds"] = wte(encoder_input_ids)
                inputs["decoder_inputs_embeds"] = wte(decoder_input_ids)

            with torch.no_grad():
                model(**inputs)[0]

    @require_torch_fp16
    def test_generate_fp16(self):
        config, input_dict = self.model_tester.prepare_config_and_inputs()
        input_ids = input_dict["input_ids"]
        attention_mask = input_ids.ne(1).to(torch_device)
        model = OPTForCausalLM(config).eval().to(torch_device)
        model.half()
        model.generate(input_ids, attention_mask=attention_mask)
        model.generate(num_beams=4, do_sample=True, early_stopping=False, num_return_sequences=3)

    def test_opt_sequence_classification_model(self):
        config, input_dict = self.model_tester.prepare_config_and_inputs()
        config.num_labels = 3
        input_ids = input_dict["input_ids"]
        attention_mask = input_ids.ne(1).to(torch_device)
        sequence_labels = ids_tensor([self.model_tester.batch_size], self.model_tester.type_sequence_label_size)
        model = OPTForSequenceClassification(config)
        model.to(torch_device)
        model.eval()
        result = model(input_ids, attention_mask=attention_mask, labels=sequence_labels)
        self.assertEqual(result.logits.shape, (self.model_tester.batch_size, self.model_tester.num_labels))

    def test_opt_sequence_classification_model_for_multi_label(self):
        config, input_dict = self.model_tester.prepare_config_and_inputs()
        config.num_labels = 3
        config.problem_type = "multi_label_classification"
        input_ids = input_dict["input_ids"]
        attention_mask = input_ids.ne(1).to(torch_device)
        sequence_labels = ids_tensor(
            [self.model_tester.batch_size, config.num_labels], self.model_tester.type_sequence_label_size
        ).to(torch.float)
        model = OPTForSequenceClassification(config)
        model.to(torch_device)
        model.eval()
        result = model(input_ids, attention_mask=attention_mask, labels=sequence_labels)
        self.assertEqual(result.logits.shape, (self.model_tester.batch_size, self.model_tester.num_labels))

    @unittest.skip("Does not work on the tiny model as we keep hitting edge cases.")
    def test_model_parallelism(self):
        super().test_model_parallelism()


def assert_tensors_close(a, b, atol=1e-12, prefix=""):
    
    if a is None and b is None:
        return True
    try:
        if torch.allclose(a, b, atol=atol):
            return True
        raise
    except Exception:
        pct_different = (torch.gt((a - b).abs(), atol)).float().mean().item()
        if a.numel() > 100:
            msg = f"tensor values are {pct_different:.1%} percent different."
        else:
            msg = f"{a} != {b}"
        if prefix:
            msg = prefix + ": " + msg
        raise AssertionError(msg)


def _long_tensor(tok_lst):
    return torch.tensor(tok_lst, dtype=torch.long, device=torch_device)


@require_torch
class OPTModelIntegrationTests(unittest.TestCase):
    @slow
    def test_inference_no_head(self):
        model = OPTModel.from_pretrained("facebook/opt-350m").to(torch_device)
        input_ids = _long_tensor([[0, 31414, 232, 328, 740, 1140, 12695, 69, 46078, 1588, 2]])

        with torch.no_grad():
            output = model(input_ids=input_ids).last_hidden_state

        expected_shape = torch.Size((1, 11, 512))
        self.assertEqual(output.shape, expected_shape)
        
        expected_slice = torch.tensor(
            [
                [-0.28726277, -1.9241608, -0.3058734],
                [-1.2737825, -0.13332152, -0.18766522],
                [0.41159445, 0.1191957, -1.3107123],
            ],
            device=torch_device,
        )
        assert_tensors_close(output[0, :3, :3], expected_slice, atol=5e-5)


@require_torch
@slow
class OPTEmbeddingsTest(unittest.TestCase):
    def setUp(self):
        super().setUp()
        self.path_model = "facebook/opt-350m"

    def test_load_model(self):
        try:
            _ = OPTForCausalLM.from_pretrained(self.path_model)
        except BaseException:
            self.fail("Failed loading model")

    def test_logits(self):
        model = OPTForCausalLM.from_pretrained(self.path_model)
        model = model.eval()
        tokenizer = GPT2Tokenizer.from_pretrained(self.path_model)

        prompts = [
            "Today is a beautiful day and I want to",
            "In the city of",
            "Paris is the capital of France and",
            "Computers and mobile phones have taken",
        ]
        
        inputs = tokenizer(prompts, return_tensors="pt", padding=True, add_special_tokens=False)
        logits = model(inputs.input_ids, attention_mask=inputs.attention_mask)[0].mean(dim=-1)
        
        logits_meta = torch.Tensor(
            [
                [1.3851, -13.8923, -10.5229, -10.7533, -0.2309, -10.2384, -0.5365, -9.0947, -5.1670],
                [-4.7073, -10.6276, -3.9415, -21.5242, -0.2822, -0.2822, -0.2822, -0.2822, -0.2822],
                [0.6247, -3.4229, -8.9179, -1.4297, -14.1650, 1.4146, -9.0218, -0.2703, -0.2703],
                [6.4783, -1.9913, -10.7926, -2.3336, 1.5092, -0.9974, -6.8213, 1.3477, 1.3477],
            ]
        )
        assert torch.allclose(logits, logits_meta, atol=1e-4)


@slow
class OPTGenerationTest(unittest.TestCase):
    @property
    def prompts(self):
        return [
            "Today is a beautiful day and I want",
            "In the city of",
            "Paris is the capital of France and",
            "Computers and mobile phones have taken",
        ]

    def test_generation_pre_attn_layer_norm(self):
        model_id = "facebook/opt-125m"

        EXPECTED_OUTPUTS = [
            "Today is a beautiful day and I want to",
            "In the city of New York, the city",
            "Paris is the capital of France and the capital",
            "Computers and mobile phones have taken over the",
        ]

        predicted_outputs = []
        tokenizer = GPT2Tokenizer.from_pretrained(model_id)
        model = OPTForCausalLM.from_pretrained(model_id)

        for prompt in self.prompts:
            input_ids = tokenizer(prompt, return_tensors="pt").input_ids

            generated_ids = model.generate(input_ids, max_length=10)

            generated_string = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
            predicted_outputs += generated_string

        self.assertListEqual(predicted_outputs, EXPECTED_OUTPUTS)

    def test_batch_generation(self):
        model_id = "facebook/opt-350m"

        tokenizer = GPT2Tokenizer.from_pretrained(model_id)
        model = OPTForCausalLM.from_pretrained(model_id)
        model.to(torch_device)

        tokenizer.padding_side = "left"

        
        sentences = [
            "Hello, my dog is a little",
            "Today, I",
        ]

        inputs = tokenizer(sentences, return_tensors="pt", padding=True)
        input_ids = inputs["input_ids"].to(torch_device)

        outputs = model.generate(
            input_ids=input_ids,
            attention_mask=inputs["attention_mask"].to(torch_device),
        )

        inputs_non_padded = tokenizer(sentences[0], return_tensors="pt").input_ids.to(torch_device)
        output_non_padded = model.generate(input_ids=inputs_non_padded)

        num_paddings = inputs_non_padded.shape[-1] - inputs["attention_mask"][-1].long().sum().cpu().item()
        inputs_padded = tokenizer(sentences[1], return_tensors="pt").input_ids.to(torch_device)
        output_padded = model.generate(input_ids=inputs_padded, max_length=model.config.max_length - num_paddings)

        batch_out_sentence = tokenizer.batch_decode(outputs, skip_special_tokens=True)
        non_padded_sentence = tokenizer.decode(output_non_padded[0], skip_special_tokens=True)
        padded_sentence = tokenizer.decode(output_padded[0], skip_special_tokens=True)

        expected_output_sentence = [
            "Hello, my dog is a little bit of a dork.\nI'm a little bit",
            "Today, I was in the middle of a conversation with a friend about the",
        ]
        self.assertListEqual(expected_output_sentence, batch_out_sentence)
        self.assertListEqual(batch_out_sentence, [non_padded_sentence, padded_sentence])

    def test_generation_post_attn_layer_norm(self):
        model_id = "facebook/opt-350m"

        EXPECTED_OUTPUTS = [
            "Today is a beautiful day and I want to",
            "In the city of San Francisco, the city",
            "Paris is the capital of France and the capital",
            "Computers and mobile phones have taken over the",
        ]

        predicted_outputs = []
        tokenizer = GPT2Tokenizer.from_pretrained(model_id)
        model = OPTForCausalLM.from_pretrained(model_id)

        for prompt in self.prompts:
            input_ids = tokenizer(prompt, return_tensors="pt").input_ids

            generated_ids = model.generate(input_ids, max_length=10)

            generated_string = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
            predicted_outputs += generated_string

        self.assertListEqual(predicted_outputs, EXPECTED_OUTPUTS)

    @require_torch_accelerator
    @require_torch_fp16
    def test_batched_nan_fp16(self):
        
        
        
        model_name = "facebook/opt-1.3b"
        tokenizer = GPT2Tokenizer.from_pretrained(model_name, use_fast=False, padding_side="left")

        model = OPTForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, use_cache=True).to(torch_device)
        model = model.eval()

        batch = tokenizer(["Who are you?", "Joe Biden is the president of"], padding=True, return_tensors="pt")

        input_ids = batch["input_ids"].to(torch_device)
        attention_mask = batch["attention_mask"].to(torch_device)

        with torch.no_grad():
            outputs = model(input_ids, attention_mask=attention_mask)
            self.assertFalse(
                torch.isnan(outputs.logits[0]).any().item()
            )  

    @slow
    def test_contrastive_search_opt(self):
        article = (
            "A chat between a curious human and the Statue of Liberty.\n\nHuman: What is your name?\nStatue: I am the "
            "Statue of Liberty.\nHuman: Where do you live?\nStatue: New York City.\nHuman: How long have you lived "
            "there?"
        )

        opt_tokenizer = GPT2Tokenizer.from_pretrained("facebook/opt-1.3b")
        opt_model = OPTForCausalLM.from_pretrained("facebook/opt-1.3b").to(torch_device)
        input_ids = opt_tokenizer(article, return_tensors="pt").input_ids.to(torch_device)

        outputs = opt_model.generate(input_ids, penalty_alpha=0.6, top_k=5, max_length=256)
        generated_text = opt_tokenizer.batch_decode(outputs, skip_special_tokens=True)

        self.assertListEqual(
            generated_text,
            [
                "A chat between a curious human and the Statue of Liberty.\n\nHuman: What is your name?\nStatue: I "
                "am the Statue of Liberty.\nHuman: Where do you live?\nStatue: New York City.\nHuman: How long have "
                "you lived there?\nStatue: A hundred years.\nHuman: And you’re from what country?\nStatue: The United "
                "States of America.\nHuman: Why did you come to America?\nStatue: I came to escape the tyranny of my "
                "country.\nHuman: What tyranny?\nStatue: They didn’t let me speak my mind.\nHuman: What was your "
                "country?\nStatue: It was a country of immigrants.\nHuman: Who were the immigrants?\nStatue: They "
                "were from all over the world.\nHuman: What language did they speak?\nStatue: French, Spanish, "
                "Italian, German, English—you name it.\nHuman: And where did they come from?\nStatue: They came from "
                "every country in the world.\nHuman: And you were born in what country?\nStatue: I was born in "
                "France.\nHuman: And your parents were French?\nStatue"
            ],
        )