


import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import YolosConfig, YolosForObjectDetection, YolosImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
_SCREAMING_SNAKE_CASE      :			Optional[Any]      				=   logging.get_logger(__name__)
def       UpperCAmelCase__						(UpperCamelCase_    ):



     """simple docstring"""






     snake_case					=    YolosConfig()

     # size of the architecture
     if "yolos_ti" in yolos_name:
          snake_case					=    1_92
          snake_case					=    7_68
          snake_case					=    12
          snake_case					=    3
          snake_case					=    [8_00, 13_33]
          snake_case					=    False
     elif yolos_name == "yolos_s_dWr":
          snake_case					=    3_30
          snake_case					=    14
          snake_case					=    6
          snake_case					=    13_20
     elif "yolos_s" in yolos_name:
          snake_case					=    3_84
          snake_case					=    15_36
          snake_case					=    12
          snake_case					=    6
     elif "yolos_b" in yolos_name:
          snake_case					=    [8_00, 13_44]

     snake_case					=    91
     snake_case					=    '''huggingface/label-files'''
     snake_case					=    '''coco-detection-id2label.json'''
     snake_case					=    json.load(open(hf_hub_download(UpperCamelCase_							,UpperCamelCase_							,repo_type='''dataset'''    )							,'''r'''    )    )
     snake_case					=    {int(UpperCamelCase_    ): v for k, v in idalabel.items()}
     snake_case					=    idalabel
     snake_case					=    {v: k for k, v in idalabel.items()}

     return config
def       UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_ = False    ):



     """simple docstring"""






     for i in range(config.num_hidden_layers    ):
          # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
          snake_case					=    state_dict.pop(F'''blocks.{i}.attn.qkv.weight'''    )
          snake_case					=    state_dict.pop(F'''blocks.{i}.attn.qkv.bias'''    )
          # next, add query, keys and values (in that order) to the state dict
          snake_case					=    in_proj_weight[: config.hidden_size, :]
          snake_case					=    in_proj_bias[: config.hidden_size]
          snake_case					=    in_proj_weight[
              config.hidden_size : config.hidden_size * 2, :
          ]
          snake_case					=    in_proj_bias[
              config.hidden_size : config.hidden_size * 2
          ]
          snake_case					=    in_proj_weight[-config.hidden_size :, :]
          snake_case					=    in_proj_bias[-config.hidden_size :]
def       UpperCAmelCase__						(UpperCamelCase_    ):



     """simple docstring"""






     if "backbone" in name:
          snake_case					=    name.replace('''backbone'''							,'''vit'''    )
     if "cls_token" in name:
          snake_case					=    name.replace('''cls_token'''							,'''embeddings.cls_token'''    )
     if "det_token" in name:
          snake_case					=    name.replace('''det_token'''							,'''embeddings.detection_tokens'''    )
     if "mid_pos_embed" in name:
          snake_case					=    name.replace('''mid_pos_embed'''							,'''encoder.mid_position_embeddings'''    )
     if "pos_embed" in name:
          snake_case					=    name.replace('''pos_embed'''							,'''embeddings.position_embeddings'''    )
     if "patch_embed.proj" in name:
          snake_case					=    name.replace('''patch_embed.proj'''							,'''embeddings.patch_embeddings.projection'''    )
     if "blocks" in name:
          snake_case					=    name.replace('''blocks'''							,'''encoder.layer'''    )
     if "attn.proj" in name:
          snake_case					=    name.replace('''attn.proj'''							,'''attention.output.dense'''    )
     if "attn" in name:
          snake_case					=    name.replace('''attn'''							,'''attention.self'''    )
     if "norm1" in name:
          snake_case					=    name.replace('''norm1'''							,'''layernorm_before'''    )
     if "norm2" in name:
          snake_case					=    name.replace('''norm2'''							,'''layernorm_after'''    )
     if "mlp.fc1" in name:
          snake_case					=    name.replace('''mlp.fc1'''							,'''intermediate.dense'''    )
     if "mlp.fc2" in name:
          snake_case					=    name.replace('''mlp.fc2'''							,'''output.dense'''    )
     if "class_embed" in name:
          snake_case					=    name.replace('''class_embed'''							,'''class_labels_classifier'''    )
     if "bbox_embed" in name:
          snake_case					=    name.replace('''bbox_embed'''							,'''bbox_predictor'''    )
     if "vit.norm" in name:
          snake_case					=    name.replace('''vit.norm'''							,'''vit.layernorm'''    )

     return name
def       UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_    ):



     """simple docstring"""






     for key in orig_state_dict.copy().keys():
          snake_case					=    orig_state_dict.pop(UpperCamelCase_    )

          if "qkv" in key:
               snake_case					=    key.split('''.'''    )
               snake_case					=    int(key_split[2]    )
               snake_case					=    model.vit.encoder.layer[layer_num].attention.attention.all_head_size
               if "weight" in key:
                    snake_case					=    val[:dim, :]
                    snake_case					=    val[
                        dim : dim * 2, :
                    ]
                    snake_case					=    val[-dim:, :]
               else:
                    snake_case					=    val[:dim]
                    snake_case					=    val[dim : dim * 2]
                    snake_case					=    val[-dim:]
          else:
               snake_case					=    val

     return orig_state_dict
def       UpperCAmelCase__						():



     """simple docstring"""






     snake_case					=    '''http://images.cocodataset.org/val2017/000000039769.jpg'''
     snake_case					=    Image.open(requests.get(UpperCamelCase_							,stream=UpperCamelCase_    ).raw    )
     return im
@torch.no_grad()
def       UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_ = False    ):



     """simple docstring"""






     snake_case					=    get_yolos_config(UpperCamelCase_    )

     # load original state_dict
     snake_case					=    torch.load(UpperCamelCase_							,map_location='''cpu'''    )['''model''']

     # load 🤗 model
     snake_case					=    YolosForObjectDetection(UpperCamelCase_    )
     model.eval()
     snake_case					=    convert_state_dict(UpperCamelCase_							,UpperCamelCase_    )
     model.load_state_dict(UpperCamelCase_    )

     # Check outputs on an image, prepared by YolosImageProcessor
     snake_case					=    8_00 if yolos_name != '''yolos_ti''' else 5_12
     snake_case					=    YolosImageProcessor(format='''coco_detection'''							,size=UpperCamelCase_    )
     snake_case					=    image_processor(images=prepare_img()							,return_tensors='''pt'''    )
     snake_case					=    model(**UpperCamelCase_    )
     snake_case			,  snake_case					=    outputs.logits, outputs.pred_boxes

     snake_case			,  snake_case					=    None, None
     if yolos_name == "yolos_ti":
          snake_case					=    torch.tensor(
              [[-39.5022, -11.9820, -17.6888], [-29.9574, -9.9769, -17.7691], [-42.3281, -20.7200, -30.6294]]    )
          snake_case					=    torch.tensor(
              [[0.4021, 0.0836, 0.7979], [0.0184, 0.2609, 0.0364], [0.1781, 0.2004, 0.2095]]    )
     elif yolos_name == "yolos_s_200_pre":
          snake_case					=    torch.tensor(
              [[-24.0248, -10.3024, -14.8290], [-42.0392, -16.8200, -27.4334], [-27.2743, -11.8154, -18.7148]]    )
          snake_case					=    torch.tensor(
              [[0.2559, 0.5455, 0.4706], [0.2989, 0.7279, 0.1875], [0.7732, 0.4017, 0.4462]]    )
     elif yolos_name == "yolos_s_300_pre":
          snake_case					=    torch.tensor(
              [[-36.2220, -14.4385, -23.5457], [-35.6970, -14.7583, -21.3935], [-31.5939, -13.6042, -16.8049]]    )
          snake_case					=    torch.tensor(
              [[0.7614, 0.2316, 0.4728], [0.7168, 0.4495, 0.3855], [0.4996, 0.1466, 0.9996]]    )
     elif yolos_name == "yolos_s_dWr":
          snake_case					=    torch.tensor(
              [[-42.8668, -24.1049, -41.1690], [-34.7456, -14.1274, -24.9194], [-33.7898, -12.1946, -25.6495]]    )
          snake_case					=    torch.tensor(
              [[0.5587, 0.2773, 0.0605], [0.5004, 0.3014, 0.9994], [0.4999, 0.1548, 0.9994]]    )
     elif yolos_name == "yolos_base":
          snake_case					=    torch.tensor(
              [[-40.6064, -24.3084, -32.6447], [-55.1990, -30.7719, -35.5877], [-51.4311, -33.3507, -35.6462]]    )
          snake_case					=    torch.tensor(
              [[0.5555, 0.2794, 0.0655], [0.9049, 0.2664, 0.1894], [0.9183, 0.1984, 0.1635]]    )
     else:
          raise ValueError(F'''Unknown yolos_name: {yolos_name}'''    )

     assert torch.allclose(logits[0, :3, :3]							,UpperCamelCase_							,atol=1e-4    )
     assert torch.allclose(pred_boxes[0, :3, :3]							,UpperCamelCase_							,atol=1e-4    )

     Path(UpperCamelCase_    ).mkdir(exist_ok=UpperCamelCase_    )
     print(F'''Saving model {yolos_name} to {pytorch_dump_folder_path}'''    )
     model.save_pretrained(UpperCamelCase_    )
     print(F'''Saving image processor to {pytorch_dump_folder_path}'''    )
     image_processor.save_pretrained(UpperCamelCase_    )

     if push_to_hub:
          snake_case					=    {
              '''yolos_ti''': '''yolos-tiny''',
              '''yolos_s_200_pre''': '''yolos-small''',
              '''yolos_s_300_pre''': '''yolos-small-300''',
              '''yolos_s_dWr''': '''yolos-small-dwr''',
              '''yolos_base''': '''yolos-base''',
          }

          print('''Pushing to the hub...'''    )
          snake_case					=    model_mapping[yolos_name]
          image_processor.push_to_hub(UpperCamelCase_							,organization='''hustvl'''    )
          model.push_to_hub(UpperCamelCase_							,organization='''hustvl'''    )


if __name__ == "__main__":
 _SCREAMING_SNAKE_CASE      :			str      				=   argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     "--yolos_name",
     default="yolos_s_200_pre",
     type=str,
     help=(
         "Name of the YOLOS model you'd like to convert. Should be one of 'yolos_ti', 'yolos_s_200_pre',"
         " 'yolos_s_300_pre', 'yolos_s_dWr', 'yolos_base'."
     ),
 )
 parser.add_argument(
     "--checkpoint_path", default=None, type=str, help="Path to the original state dict (.pth file)."
 )
 parser.add_argument(
     "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
 )
 parser.add_argument(
     "--push_to_hub", action="store_true", help="Whether or not to push the converted model to the 🤗 hub."
 )

 _SCREAMING_SNAKE_CASE      :			List[str]      				=   parser.parse_args()
 convert_yolos_checkpoint(args.yolos_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub)








import unittest

from parameterized import parameterized

from transformers import AutoTokenizer, GPTNeoXConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
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 (
     GPTNeoXForCausalLM,
     GPTNeoXForQuestionAnswering,
     GPTNeoXForSequenceClassification,
     GPTNeoXForTokenClassification,
     GPTNeoXModel,
 )

class   A__ :







    """simple docstring"""

    def __init__(	self      , __snake_case      , __snake_case=1_3      , __snake_case=7      , __snake_case=True      , __snake_case=True      , __snake_case=True      , __snake_case=True      , __snake_case=9_9      , __snake_case=6_4      , __snake_case=5      , __snake_case=4      , __snake_case=3_7      , __snake_case="gelu"      , __snake_case=0.1      , __snake_case=0.1      , __snake_case=5_1_2      , __snake_case=1_6      , __snake_case=2      , __snake_case=0.02      , __snake_case=3      , __snake_case=4      , __snake_case=None      , ):
         snake_case					=    parent
         snake_case					=    batch_size
         snake_case					=    seq_length
         snake_case					=    is_training
         snake_case					=    use_input_mask
         snake_case					=    use_token_type_ids
         snake_case					=    use_labels
         snake_case					=    vocab_size
         snake_case					=    hidden_size
         snake_case					=    num_hidden_layers
         snake_case					=    num_attention_heads
         snake_case					=    intermediate_size
         snake_case					=    hidden_act
         snake_case					=    hidden_dropout_prob
         snake_case					=    attention_probs_dropout_prob
         snake_case					=    max_position_embeddings
         snake_case					=    type_vocab_size
         snake_case					=    type_sequence_label_size
         snake_case					=    initializer_range
         snake_case					=    num_labels
         snake_case					=    num_choices
         snake_case					=    scope
         snake_case					=    vocab_size - 1
    def 		a_      (	self   ):
         snake_case					=    ids_tensor([self.batch_size, self.seq_length]      , self.vocab_size   )

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

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

         snake_case					=    self.get_config()

         return config, input_ids, input_mask, token_labels
    def 		a_      (	self   ):
         return GPTNeoXConfig(
             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_size=self.intermediate_size      , hidden_act=self.hidden_act      , hidden_dropout_prob=self.hidden_dropout_prob      , attention_probs_dropout_prob=self.attention_probs_dropout_prob      , max_position_embeddings=self.max_position_embeddings      , type_vocab_size=self.type_vocab_size      , is_decoder=__snake_case      , initializer_range=self.initializer_range      , pad_token_id=self.pad_token_id      , )
    def 		a_      (	self   ):
         snake_case			,  snake_case			,  snake_case			,  snake_case					=    self.prepare_config_and_inputs()

         snake_case					=    True

         return config, input_ids, input_mask, token_labels
    def 		a_      (	self      , __snake_case      , __snake_case      , __snake_case   ):
         snake_case					=    GPTNeoXModel(config=__snake_case   )
         model.to(__snake_case   )
         model.eval()
         snake_case					=    model(__snake_case      , attention_mask=__snake_case   )
         snake_case					=    model(__snake_case   )
         self.parent.assertEqual(result.last_hidden_state.shape      , (self.batch_size, self.seq_length, self.hidden_size)   )
    def 		a_      (	self      , __snake_case      , __snake_case      , __snake_case   ):
         snake_case					=    True
         snake_case					=    GPTNeoXModel(__snake_case   )
         model.to(__snake_case   )
         model.eval()
         snake_case					=    model(__snake_case      , attention_mask=__snake_case   )
         self.parent.assertEqual(result.last_hidden_state.shape      , (self.batch_size, self.seq_length, self.hidden_size)   )
    def 		a_      (	self      , __snake_case      , __snake_case      , __snake_case      , __snake_case   ):
         snake_case					=    GPTNeoXForCausalLM(config=__snake_case   )
         model.to(__snake_case   )
         model.eval()
         snake_case					=    model(__snake_case      , attention_mask=__snake_case      , labels=__snake_case   )
         self.parent.assertEqual(result.logits.shape      , (self.batch_size, self.seq_length, self.vocab_size)   )
    def 		a_      (	self      , __snake_case      , __snake_case      , __snake_case      , __snake_case   ):
         snake_case					=    self.num_labels
         snake_case					=    GPTNeoXForQuestionAnswering(__snake_case   )
         model.to(__snake_case   )
         model.eval()
         snake_case					=    model(__snake_case      , attention_mask=__snake_case   )
         self.parent.assertEqual(result.start_logits.shape      , (self.batch_size, self.seq_length)   )
         self.parent.assertEqual(result.end_logits.shape      , (self.batch_size, self.seq_length)   )
    def 		a_      (	self      , __snake_case      , __snake_case      , __snake_case      , __snake_case   ):
         snake_case					=    self.num_labels
         snake_case					=    GPTNeoXForSequenceClassification(__snake_case   )
         model.to(__snake_case   )
         model.eval()
         snake_case					=    ids_tensor([self.batch_size]      , self.type_sequence_label_size   )
         snake_case					=    model(__snake_case      , attention_mask=__snake_case      , labels=__snake_case   )
         self.parent.assertEqual(result.logits.shape      , (self.batch_size, self.num_labels)   )
    def 		a_      (	self      , __snake_case      , __snake_case      , __snake_case      , __snake_case   ):
         snake_case					=    self.num_labels
         snake_case					=    GPTNeoXForTokenClassification(__snake_case   )
         model.to(__snake_case   )
         model.eval()
         snake_case					=    model(__snake_case      , attention_mask=__snake_case      , labels=__snake_case   )
         self.parent.assertEqual(result.logits.shape      , (self.batch_size, self.seq_length, self.num_labels)   )
    def 		a_      (	self      , __snake_case      , __snake_case      , __snake_case   ):
         snake_case					=    True
         snake_case					=    GPTNeoXForCausalLM(config=__snake_case   )
         model.to(__snake_case   )
         model.eval()

         # first forward pass
         snake_case					=    model(__snake_case      , attention_mask=__snake_case      , use_cache=__snake_case   )
         snake_case					=    outputs.past_key_values

         # create hypothetical multiple next token and extent to next_input_ids
         snake_case					=    ids_tensor((self.batch_size, 3)      , config.vocab_size   )
         snake_case					=    ids_tensor((self.batch_size, 3)      , vocab_size=2   )

         # append to next input_ids and
         snake_case					=    torch.cat([input_ids, next_tokens]      , dim=-1   )
         snake_case					=    torch.cat([input_mask, next_mask]      , dim=-1   )

         snake_case					=    model(__snake_case      , attention_mask=__snake_case      , output_hidden_states=__snake_case   )
         snake_case					=    output_from_no_past['''hidden_states'''][0]
         snake_case					=    model(
             __snake_case      , attention_mask=__snake_case      , past_key_values=__snake_case      , output_hidden_states=__snake_case      , )['''hidden_states'''][0]

         # select random slice
         snake_case					=    ids_tensor((1,)      , output_from_past.shape[-1]   ).item()
         snake_case					=    output_from_no_past[:, -3:, random_slice_idx].detach()
         snake_case					=    output_from_past[:, :, random_slice_idx].detach()

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

         # test that outputs are equal for slice
         self.parent.assertTrue(torch.allclose(__snake_case      , __snake_case      , atol=1E-3   )   )




    def 		a_      (	self   ):
         snake_case					=    self.prepare_config_and_inputs()
         snake_case			,  snake_case			,  snake_case			,  snake_case					=    config_and_inputs
         snake_case					=    {'''input_ids''': input_ids, '''attention_mask''': input_mask}
         return config, inputs_dict






@require_torch
class   A__ (   snake_case__	, snake_case__	, snake_case__	, unittest.TestCase     ):







    """simple docstring"""

    __magic_name__        =		(
        (
            GPTNeoXModel,
            GPTNeoXForCausalLM,
            GPTNeoXForQuestionAnswering,
            GPTNeoXForSequenceClassification,
            GPTNeoXForTokenClassification,
        )
        if is_torch_available()
        else ()
    )
    __magic_name__        =		(GPTNeoXForCausalLM,) if is_torch_available() else ()
    __magic_name__        =		(
        {
            'feature-extraction': GPTNeoXModel,
            'question-answering': GPTNeoXForQuestionAnswering,
            'text-classification': GPTNeoXForSequenceClassification,
            'text-generation': GPTNeoXForCausalLM,
            'token-classification': GPTNeoXForTokenClassification,
            'zero-shot': GPTNeoXForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    __magic_name__        =		False
    __magic_name__        =		False
    __magic_name__        =		False
    __magic_name__        =		False
    def 		a_      (	self   ):
         snake_case					=    GPTNeoXModelTester(self   )
         snake_case					=    ConfigTester(self      , config_class=__snake_case      , hidden_size=6_4      , num_attention_heads=8   )
    def 		a_      (	self   ):
         self.config_tester.run_common_tests()
    def 		a_      (	self   ):
         snake_case			,  snake_case			,  snake_case			,  snake_case					=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(__snake_case      , __snake_case      , __snake_case   )
    def 		a_      (	self   ):
         snake_case			,  snake_case			,  snake_case			,  snake_case					=    self.model_tester.prepare_config_and_inputs_for_decoder()
         self.model_tester.create_and_check_model_as_decoder(__snake_case      , __snake_case      , __snake_case   )
    def 		a_      (	self   ):
         # This regression test was failing with PyTorch < 1.3
         snake_case			,  snake_case			,  snake_case			,  snake_case					=    self.model_tester.prepare_config_and_inputs_for_decoder()

         snake_case					=    None

         self.model_tester.create_and_check_model_as_decoder(__snake_case      , __snake_case      , __snake_case   )
    def 		a_      (	self   ):
         snake_case			,  snake_case			,  snake_case			,  snake_case					=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_decoder_model_past_large_inputs(__snake_case      , __snake_case      , __snake_case   )
    def 		a_      (	self   ):
         snake_case					=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_causal_lm(*__snake_case   )
    def 		a_      (	self   ):
         snake_case					=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_question_answering(*__snake_case   )
    def 		a_      (	self   ):
         snake_case					=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_sequence_classification(*__snake_case   )
    def 		a_      (	self   ):
         snake_case					=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(*__snake_case   )
    @unittest.skip(reason='''Feed forward chunking is not implemented'''   )
    def 		a_      (	self   ):
         pass




    @parameterized.expand([('''linear''',), ('''dynamic''',)]   )
    def 		a_      (	self      , __snake_case   ):
         snake_case			,  snake_case					=    self.model_tester.prepare_config_and_inputs_for_common()
         snake_case					=    ids_tensor([1, 1_0]      , config.vocab_size   )
         snake_case					=    ids_tensor([1, int(config.max_position_embeddings * 1.5   )]      , config.vocab_size   )

         set_seed(4_2   )  # Fixed seed at init time so the two models get the same random weights
         snake_case					=    GPTNeoXModel(__snake_case   )
         original_model.to(__snake_case   )
         original_model.eval()
         snake_case					=    original_model(__snake_case   ).last_hidden_state
         snake_case					=    original_model(__snake_case   ).last_hidden_state

         set_seed(4_2   )  # Fixed seed at init time so the two models get the same random weights
         snake_case					=    {'''type''': scaling_type, '''factor''': 10.0}
         snake_case					=    GPTNeoXModel(__snake_case   )
         scaled_model.to(__snake_case   )
         scaled_model.eval()
         snake_case					=    scaled_model(__snake_case   ).last_hidden_state
         snake_case					=    scaled_model(__snake_case   ).last_hidden_state

         # Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
         # maximum sequence length, so the outputs for the short input should match.
         if scaling_type == "dynamic":
              self.assertTrue(torch.allclose(__snake_case      , __snake_case      , atol=1E-5   )   )
         else:
              self.assertFalse(torch.allclose(__snake_case      , __snake_case      , atol=1E-5   )   )

         # The output should be different for long inputs
         self.assertFalse(torch.allclose(__snake_case      , __snake_case      , atol=1E-5   )   )






@require_torch
class   A__ (   unittest.TestCase     ):







    """simple docstring"""

    @slow
    def 		a_      (	self   ):
         snake_case					=    AutoTokenizer.from_pretrained('''EleutherAI/pythia-410m-deduped'''   )
         for checkpointing in [True, False]:
              snake_case					=    GPTNeoXForCausalLM.from_pretrained('''EleutherAI/pythia-410m-deduped'''   )

              if checkpointing:
                   model.gradient_checkpointing_enable()
              else:
                   model.gradient_checkpointing_disable()
              model.to(__snake_case   )

              snake_case					=    tokenizer('''My favorite food is'''      , return_tensors='''pt'''   ).to(__snake_case   )
              # The hub repo. is updated on 2023-04-04, resulting in poor outputs.
              # See: https://github.com/huggingface/transformers/pull/24193
              snake_case					=    '''My favorite food is a good old-fashioned, old-fashioned, old-fashioned.\n\nI\'m not sure'''

              snake_case					=    model.generate(**__snake_case      , do_sample=__snake_case      , max_new_tokens=2_0   )
              snake_case					=    tokenizer.batch_decode(__snake_case   )[0]

              self.assertEqual(__snake_case      , __snake_case   )




