





import argparse
import json
from dataclasses import dataclass, field
from functools import partial
from pathlib import Path
from typing import Callable, Dict, List, Tuple

import timm
import torch
import torch.nn as nn
from classy_vision.models.regnet import RegNet, RegNetParams, RegNetYaagf, RegNetYaagf, RegNetYaaagf
from huggingface_hub import cached_download, hf_hub_url
from torch import Tensor
from vissl.models.model_helpers import get_trunk_forward_outputs

from transformers import AutoImageProcessor, RegNetConfig, RegNetForImageClassification, RegNetModel
from transformers.utils import logging


logging.set_verbosity_info()
_UpperCAmelCase							:			Optional[int]        =							logging.get_logger()

@dataclass
class   lowerCAmelCase				:
 UpperCAmelCase__						=						42
 UpperCAmelCase__						=						field(default_factory=__UpperCamelCase     )
 UpperCAmelCase__						=						field(default_factory=__UpperCamelCase     )
 def 	A_      ( self    : Union[str, Any]     ,	UpperCAmelCase    : int     ,	UpperCAmelCase    : Tensor     ,	UpperCAmelCase    : Tensor      )     ->      Union[str, Any]:
     lowerCamelCase__       : str   					=						len(list(m.modules()      )      ) == 1 or isinstance(UpperCAmelCase     ,	nn.Convad      ) or isinstance(UpperCAmelCase     ,	nn.BatchNormad      )
     if has_not_submodules:
         self.traced.append(UpperCAmelCase      )
 def __call__( self    : Optional[int]     ,	UpperCAmelCase    : Tensor      )     ->      List[Any]:
     for m in self.module.modules():
         self.handles.append(m.register_forward_hook(self._forward_hook      )      )
     self.module(UpperCAmelCase      )
     [x.remove() for x in self.handles]
     return self







 @property
 def 	A_      ( self    : Tuple      )     ->      Union[str, Any]:
     # check the len of the state_dict keys to see if we have learnable params
     return list(filter(lambda UpperCAmelCase      : len(list(x.state_dict().keys()      )      ) > 0     ,	self.traced      )      )


@dataclass
class   lowerCAmelCase				:
 UpperCAmelCase__						=						42
 UpperCAmelCase__						=						42
 UpperCAmelCase__						=						1
 UpperCAmelCase__						=						field(default_factory=__UpperCamelCase     )
 UpperCAmelCase__						=						field(default_factory=__UpperCamelCase     )
 UpperCAmelCase__						=						True
 def __call__( self    : Any     ,	UpperCAmelCase    : Tensor      )     ->      Any:
     lowerCamelCase__       : List[str]   					=						Tracker(self.dest      )(UpperCAmelCase      ).parametrized
     lowerCamelCase__       : int   					=						Tracker(self.src      )(UpperCAmelCase      ).parametrized

     lowerCamelCase__       : int   					=						list(filter(lambda UpperCAmelCase      : type(UpperCAmelCase      ) not in self.src_skip     ,	UpperCAmelCase      )      )
     lowerCamelCase__       : Union[str, Any]   					=						list(filter(lambda UpperCAmelCase      : type(UpperCAmelCase      ) not in self.dest_skip     ,	UpperCAmelCase      )      )

     if len(UpperCAmelCase      ) != len(UpperCAmelCase      ) and self.raise_if_mismatch:
         raise Exception(
             F"""Numbers of operations are different. Source module has {len(UpperCAmelCase      )} operations while"""
             F""" destination module has {len(UpperCAmelCase      )}."""      )

     for dest_m, src_m in zip(UpperCAmelCase     ,	UpperCAmelCase      ):
         dest_m.load_state_dict(src_m.state_dict()      )
         if self.verbose == 1:
             print(F"""Transfered from={src_m} to={dest_m}"""      )


class   lowerCAmelCase				(   nn.Module     ):
 def __init__( self    : Any     ,	UpperCAmelCase    : nn.Module      )     ->      Union[str, Any]:
     super().__init__()

     lowerCamelCase__       : List[Tuple[str, nn.Module]]   					=						[]
     # - get the stem
     feature_blocks.append(('conv1', model.stem)      )
     # - get all the feature blocks
     for k, v in model.trunk_output.named_children():
         assert k.startswith('block'      ), F"""Unexpected layer name {k}"""
         lowerCamelCase__       : Union[str, Any]   					=						len(UpperCAmelCase      ) + 1
         feature_blocks.append((F"""res{block_index}""", v)      )

     lowerCamelCase__       : int   					=						nn.ModuleDict(UpperCAmelCase      )







 def 	A_      ( self    : List[str]     ,	UpperCAmelCase    : Tensor      )     ->      Optional[Any]:
     return get_trunk_forward_outputs(
         UpperCAmelCase     ,	out_feat_keys=UpperCAmelCase     ,	feature_blocks=self._feature_blocks     ,	)


class   lowerCAmelCase				(   __UpperCamelCase     ):
 def 	A_      ( self    : Dict     ,	UpperCAmelCase    : str      )     ->      str:
     lowerCamelCase__       : Any   					=						x.split('-'      )
     return x_split[0] + x_split[1] + "_" + "".join(x_split[2:]      )







 def __getitem__( self    : Tuple     ,	UpperCAmelCase    : str      )     ->      Callable[[], Tuple[nn.Module, Dict]]:
     # default to timm!
     if x not in self:
         lowerCamelCase__       : Optional[int]   					=						self.convert_name_to_timm(UpperCAmelCase      )
         lowerCamelCase__       : str   					=						partial(lambda: (timm.create_model(UpperCAmelCase     ,	pretrained=UpperCAmelCase      ).eval(), None)      )

     else:
         lowerCamelCase__       : Dict   					=						super().__getitem__(UpperCAmelCase      )

     return val


class   lowerCAmelCase				(   __UpperCamelCase     ):
 def __getitem__( self    : Optional[Any]     ,	UpperCAmelCase    : str      )     ->      Callable[[], nn.Module]:
     if "seer" in x and "in1k" not in x:
         lowerCamelCase__       : str   					=						RegNetModel
     else:
         lowerCamelCase__       : Union[str, Any]   					=						RegNetForImageClassification
     return val


def 			SCREAMING_SNAKE_CASE (			_UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase					)     ->	List[Any]:
    for from_key, to_key in keys:
        lowerCamelCase__       : str   					=						from_state_dict[from_key].clone()
        print(F"""Copied key={from_key} to={to_key}"""					)
    return to_state_dict


def 			SCREAMING_SNAKE_CASE (			_UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase = True		,      )     ->	List[str]:
    print(F"""Converting {name}..."""					)
    with torch.no_grad():
        lowerCamelCase__					,				lowerCamelCase__       : Optional[int]   					=						from_model_func()
        lowerCamelCase__       : Dict   					=						our_model_func(_UpperCAmelCase					).eval()
        lowerCamelCase__       : Optional[Any]   					=						ModuleTransfer(src=_UpperCAmelCase		,      dest=_UpperCAmelCase		,      raise_if_mismatch=_UpperCAmelCase					)
        lowerCamelCase__       : int   					=						torch.randn((1, 3, 224, 224)					)
        module_transfer(_UpperCAmelCase					)

    if from_state_dict is not None:
        lowerCamelCase__       : Dict   					=						[]
        # for seer - in1k finetuned we have to manually copy the head
        if "seer" in name and "in1k" in name:
            lowerCamelCase__       : Optional[Any]   					=						[('0.clf.0.weight', 'classifier.1.weight'), ('0.clf.0.bias', 'classifier.1.bias')]
        lowerCamelCase__       : Tuple   					=						manually_copy_vissl_head(_UpperCAmelCase		,      our_model.state_dict()		,      _UpperCAmelCase					)
        our_model.load_state_dict(_UpperCAmelCase					)

    lowerCamelCase__       : Any   					=						our_model(_UpperCAmelCase		,      output_hidden_states=_UpperCAmelCase					)
    lowerCamelCase__       : Dict   					=						(
        our_outputs.logits if isinstance(_UpperCAmelCase		,      _UpperCAmelCase					) else our_outputs.last_hidden_state
    )

    lowerCamelCase__       : List[str]   					=						from_model(_UpperCAmelCase					)
    lowerCamelCase__       : Optional[Any]   					=						from_output[-1] if type(_UpperCAmelCase					) is list else from_output

    # now since I don't want to use any config files, vissl seer model doesn't actually have an head, so let's just check the last hidden state
    if "seer" in name and "in1k" in name:
        lowerCamelCase__       : int   					=						our_outputs.hidden_states[-1]

    assert torch.allclose(_UpperCAmelCase		,      _UpperCAmelCase					), "The model logits don't match the original one."

    if push_to_hub:
        our_model.push_to_hub(
            repo_path_or_name=save_directory / name		,      commit_message='Add model'		,      use_temp_dir=_UpperCAmelCase		,      )

        lowerCamelCase__       : List[Any]   					=						224 if 'seer' not in name else 384
        # we can use the convnext one
        lowerCamelCase__       : Dict   					=						AutoImageProcessor.from_pretrained('facebook/convnext-base-224-22k-1k'		,      size=_UpperCAmelCase					)
        image_processor.push_to_hub(
            repo_path_or_name=save_directory / name		,      commit_message='Add image processor'		,      use_temp_dir=_UpperCAmelCase		,      )

        print(F"""Pushed {name}"""					)


def 			SCREAMING_SNAKE_CASE (			_UpperCAmelCase		,      _UpperCAmelCase = None		,      _UpperCAmelCase = True					)     ->	str:
    lowerCamelCase__       : List[Any]   					=						'imagenet-1k-id2label.json'
    lowerCamelCase__       : Any   					=						1000
    lowerCamelCase__       : Dict   					=						(1, num_labels)

    lowerCamelCase__       : List[str]   					=						'huggingface/label-files'
    lowerCamelCase__       : Tuple   					=						num_labels
    lowerCamelCase__       : List[str]   					=						json.load(open(cached_download(hf_hub_url(_UpperCAmelCase		,      _UpperCAmelCase		,      repo_type='dataset'					)					)		,      'r'					)					)
    lowerCamelCase__       : Union[str, Any]   					=						{int(_UpperCAmelCase					): v for k, v in idalabel.items()}

    lowerCamelCase__       : Union[str, Any]   					=						idalabel
    lowerCamelCase__       : Tuple   					=						{v: k for k, v in idalabel.items()}

    lowerCamelCase__       : Optional[int]   					=						partial(_UpperCAmelCase		,      num_labels=_UpperCAmelCase		,      idalabel=_UpperCAmelCase		,      labelaid=_UpperCAmelCase					)

    lowerCamelCase__       : Optional[Any]   					=						{
        'regnet-x-002': ImageNetPreTrainedConfig(
            depths=[1, 1, 4, 7]		,      hidden_sizes=[24, 56, 152, 368]		,      groups_width=8		,      layer_type='x'					),
        'regnet-x-004': ImageNetPreTrainedConfig(
            depths=[1, 2, 7, 12]		,      hidden_sizes=[32, 64, 160, 384]		,      groups_width=16		,      layer_type='x'					),
        'regnet-x-006': ImageNetPreTrainedConfig(
            depths=[1, 3, 5, 7]		,      hidden_sizes=[48, 96, 240, 528]		,      groups_width=24		,      layer_type='x'					),
        'regnet-x-008': ImageNetPreTrainedConfig(
            depths=[1, 3, 7, 5]		,      hidden_sizes=[64, 128, 288, 672]		,      groups_width=16		,      layer_type='x'					),
        'regnet-x-016': ImageNetPreTrainedConfig(
            depths=[2, 4, 10, 2]		,      hidden_sizes=[72, 168, 408, 912]		,      groups_width=24		,      layer_type='x'					),
        'regnet-x-032': ImageNetPreTrainedConfig(
            depths=[2, 6, 15, 2]		,      hidden_sizes=[96, 192, 432, 1008]		,      groups_width=48		,      layer_type='x'					),
        'regnet-x-040': ImageNetPreTrainedConfig(
            depths=[2, 5, 14, 2]		,      hidden_sizes=[80, 240, 560, 1360]		,      groups_width=40		,      layer_type='x'					),
        'regnet-x-064': ImageNetPreTrainedConfig(
            depths=[2, 4, 10, 1]		,      hidden_sizes=[168, 392, 784, 1624]		,      groups_width=56		,      layer_type='x'					),
        'regnet-x-080': ImageNetPreTrainedConfig(
            depths=[2, 5, 15, 1]		,      hidden_sizes=[80, 240, 720, 1920]		,      groups_width=120		,      layer_type='x'					),
        'regnet-x-120': ImageNetPreTrainedConfig(
            depths=[2, 5, 11, 1]		,      hidden_sizes=[224, 448, 896, 2240]		,      groups_width=112		,      layer_type='x'					),
        'regnet-x-160': ImageNetPreTrainedConfig(
            depths=[2, 6, 13, 1]		,      hidden_sizes=[256, 512, 896, 2048]		,      groups_width=128		,      layer_type='x'					),
        'regnet-x-320': ImageNetPreTrainedConfig(
            depths=[2, 7, 13, 1]		,      hidden_sizes=[336, 672, 1344, 2520]		,      groups_width=168		,      layer_type='x'					),
        # y variant
        'regnet-y-002': ImageNetPreTrainedConfig(depths=[1, 1, 4, 7]		,      hidden_sizes=[24, 56, 152, 368]		,      groups_width=8					),
        'regnet-y-004': ImageNetPreTrainedConfig(
            depths=[1, 3, 6, 6]		,      hidden_sizes=[48, 104, 208, 440]		,      groups_width=8					),
        'regnet-y-006': ImageNetPreTrainedConfig(
            depths=[1, 3, 7, 4]		,      hidden_sizes=[48, 112, 256, 608]		,      groups_width=16					),
        'regnet-y-008': ImageNetPreTrainedConfig(
            depths=[1, 3, 8, 2]		,      hidden_sizes=[64, 128, 320, 768]		,      groups_width=16					),
        'regnet-y-016': ImageNetPreTrainedConfig(
            depths=[2, 6, 17, 2]		,      hidden_sizes=[48, 120, 336, 888]		,      groups_width=24					),
        'regnet-y-032': ImageNetPreTrainedConfig(
            depths=[2, 5, 13, 1]		,      hidden_sizes=[72, 216, 576, 1512]		,      groups_width=24					),
        'regnet-y-040': ImageNetPreTrainedConfig(
            depths=[2, 6, 12, 2]		,      hidden_sizes=[128, 192, 512, 1088]		,      groups_width=64					),
        'regnet-y-064': ImageNetPreTrainedConfig(
            depths=[2, 7, 14, 2]		,      hidden_sizes=[144, 288, 576, 1296]		,      groups_width=72					),
        'regnet-y-080': ImageNetPreTrainedConfig(
            depths=[2, 4, 10, 1]		,      hidden_sizes=[168, 448, 896, 2016]		,      groups_width=56					),
        'regnet-y-120': ImageNetPreTrainedConfig(
            depths=[2, 5, 11, 1]		,      hidden_sizes=[224, 448, 896, 2240]		,      groups_width=112					),
        'regnet-y-160': ImageNetPreTrainedConfig(
            depths=[2, 4, 11, 1]		,      hidden_sizes=[224, 448, 1232, 3024]		,      groups_width=112					),
        'regnet-y-320': ImageNetPreTrainedConfig(
            depths=[2, 5, 12, 1]		,      hidden_sizes=[232, 696, 1392, 3712]		,      groups_width=232					),
        # models created by SEER -> https://arxiv.org/abs/2202.08360
        'regnet-y-320-seer': RegNetConfig(depths=[2, 5, 12, 1]		,      hidden_sizes=[232, 696, 1392, 3712]		,      groups_width=232					),
        'regnet-y-640-seer': RegNetConfig(depths=[2, 5, 12, 1]		,      hidden_sizes=[328, 984, 1968, 4920]		,      groups_width=328					),
        'regnet-y-1280-seer': RegNetConfig(
            depths=[2, 7, 17, 1]		,      hidden_sizes=[528, 1056, 2904, 7392]		,      groups_width=264					),
        'regnet-y-2560-seer': RegNetConfig(
            depths=[3, 7, 16, 1]		,      hidden_sizes=[640, 1696, 2544, 5088]		,      groups_width=640					),
        'regnet-y-10b-seer': ImageNetPreTrainedConfig(
            depths=[2, 7, 17, 1]		,      hidden_sizes=[2020, 4040, 1_1110, 2_8280]		,      groups_width=1010					),
        # finetuned on imagenet
        'regnet-y-320-seer-in1k': ImageNetPreTrainedConfig(
            depths=[2, 5, 12, 1]		,      hidden_sizes=[232, 696, 1392, 3712]		,      groups_width=232					),
        'regnet-y-640-seer-in1k': ImageNetPreTrainedConfig(
            depths=[2, 5, 12, 1]		,      hidden_sizes=[328, 984, 1968, 4920]		,      groups_width=328					),
        'regnet-y-1280-seer-in1k': ImageNetPreTrainedConfig(
            depths=[2, 7, 17, 1]		,      hidden_sizes=[528, 1056, 2904, 7392]		,      groups_width=264					),
        'regnet-y-2560-seer-in1k': ImageNetPreTrainedConfig(
            depths=[3, 7, 16, 1]		,      hidden_sizes=[640, 1696, 2544, 5088]		,      groups_width=640					),
        'regnet-y-10b-seer-in1k': ImageNetPreTrainedConfig(
            depths=[2, 7, 17, 1]		,      hidden_sizes=[2020, 4040, 1_1110, 2_8280]		,      groups_width=1010					),
    }

    lowerCamelCase__       : Optional[Any]   					=						NameToOurModelFuncMap()
    lowerCamelCase__       : Optional[Any]   					=						NameToFromModelFuncMap()
    # add seer weights logic

    def load_using_classy_vision(_UpperCAmelCase		,      _UpperCAmelCase					) -> Tuple[nn.Module, Dict]:
        lowerCamelCase__       : List[Any]   					=						torch.hub.load_state_dict_from_url(_UpperCAmelCase		,      model_dir=str(_UpperCAmelCase					)		,      map_location='cpu'					)
        lowerCamelCase__       : int   					=						model_func()
        # check if we have a head, if yes add it
        lowerCamelCase__       : Optional[int]   					=						files['classy_state_dict']['base_model']['model']
        lowerCamelCase__       : Any   					=						model_state_dict['trunk']
        model.load_state_dict(_UpperCAmelCase					)
        return model.eval(), model_state_dict["heads"]

    # pretrained
    lowerCamelCase__       : Union[str, Any]   					=						partial(
        _UpperCAmelCase		,      'https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_regnet32d/seer_regnet32gf_model_iteration244000.torch'		,      lambda: FakeRegNetVisslWrapper(RegNetYaagf()					)		,      )

    lowerCamelCase__       : str   					=						partial(
        _UpperCAmelCase		,      'https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_regnet64/seer_regnet64gf_model_final_checkpoint_phase0.torch'		,      lambda: FakeRegNetVisslWrapper(RegNetYaagf()					)		,      )

    lowerCamelCase__       : Tuple   					=						partial(
        _UpperCAmelCase		,      'https://dl.fbaipublicfiles.com/vissl/model_zoo/swav_ig1b_regnet128Gf_cnstant_bs32_node16_sinkhorn10_proto16k_syncBN64_warmup8k/model_final_checkpoint_phase0.torch'		,      lambda: FakeRegNetVisslWrapper(RegNetYaaagf()					)		,      )

    lowerCamelCase__       : Any   					=						partial(
        _UpperCAmelCase		,      'https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_regnet10B/model_iteration124500_conso.torch'		,      lambda: FakeRegNetVisslWrapper(
            RegNet(RegNetParams(depth=27		,      group_width=1010		,      w_a=1744		,      w_a=620.83		,      w_m=2.52					)					)					)		,      )

    # IN1K finetuned
    lowerCamelCase__       : List[Any]   					=						partial(
        _UpperCAmelCase		,      'https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_finetuned/seer_regnet32_finetuned_in1k_model_final_checkpoint_phase78.torch'		,      lambda: FakeRegNetVisslWrapper(RegNetYaagf()					)		,      )

    lowerCamelCase__       : str   					=						partial(
        _UpperCAmelCase		,      'https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_finetuned/seer_regnet64_finetuned_in1k_model_final_checkpoint_phase78.torch'		,      lambda: FakeRegNetVisslWrapper(RegNetYaagf()					)		,      )

    lowerCamelCase__       : Any   					=						partial(
        _UpperCAmelCase		,      'https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_finetuned/seer_regnet128_finetuned_in1k_model_final_checkpoint_phase78.torch'		,      lambda: FakeRegNetVisslWrapper(RegNetYaaagf()					)		,      )

    lowerCamelCase__       : Any   					=						partial(
        _UpperCAmelCase		,      'https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_finetuned/seer_10b_finetuned_in1k_model_phase28_conso.torch'		,      lambda: FakeRegNetVisslWrapper(
            RegNet(RegNetParams(depth=27		,      group_width=1010		,      w_a=1744		,      w_a=620.83		,      w_m=2.52					)					)					)		,      )

    if model_name:
        convert_weight_and_push(
            _UpperCAmelCase		,      names_to_from_model_map[model_name]		,      names_to_ours_model_map[model_name]		,      names_to_config[model_name]		,      _UpperCAmelCase		,      _UpperCAmelCase		,      )
    else:
        for model_name, config in names_to_config.items():
            convert_weight_and_push(
                _UpperCAmelCase		,      names_to_from_model_map[model_name]		,      names_to_ours_model_map[model_name]		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      )
    return config, expected_shape


if __name__ == "__main__":
    _UpperCAmelCase							:			Optional[int]        =							argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        """--model_name""",
        default=None,
        type=str,
        help=(
            """The name of the model you wish to convert, it must be one of the supported regnet* architecture,"""
            """ currently: regnetx-*, regnety-*. If `None`, all of them will the converted."""
        ),
    )
    parser.add_argument(
        """--pytorch_dump_folder_path""",
        default=None,
        type=Path,
        required=True,
        help="""Path to the output PyTorch model directory.""",
    )
    parser.add_argument(
        """--push_to_hub""",
        default=True,
        type=bool,
        required=False,
        help="""If True, push model and image processor to the hub.""",
    )

    _UpperCAmelCase							:			Any        =							parser.parse_args()

    _UpperCAmelCase							:			Path        =							args.pytorch_dump_folder_path
    pytorch_dump_folder_path.mkdir(exist_ok=True, parents=True)
    convert_weights_and_push(pytorch_dump_folder_path, args.model_name, args.push_to_hub)












import unittest

from transformers import AutoTokenizer, FalconConfig, is_torch_available
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 (
        FalconForCausalLM,
        FalconForQuestionAnswering,
        FalconForSequenceClassification,
        FalconForTokenClassification,
        FalconModel,
    )

class   lowerCAmelCase				:
 def __init__( self    : List[Any]     ,	UpperCAmelCase    : int     ,	UpperCAmelCase    : Any=3     ,	UpperCAmelCase    : List[str]=7     ,	UpperCAmelCase    : Any=True     ,	UpperCAmelCase    : str=True     ,	UpperCAmelCase    : Dict=False     ,	UpperCAmelCase    : Union[str, Any]=True     ,	UpperCAmelCase    : int=99     ,	UpperCAmelCase    : Tuple=32     ,	UpperCAmelCase    : int=5     ,	UpperCAmelCase    : str=4     ,	UpperCAmelCase    : Optional[Any]=37     ,	UpperCAmelCase    : Dict="gelu"     ,	UpperCAmelCase    : Optional[int]=0.1     ,	UpperCAmelCase    : Any=0.1     ,	UpperCAmelCase    : List[str]=512     ,	UpperCAmelCase    : Dict=16     ,	UpperCAmelCase    : Optional[int]=2     ,	UpperCAmelCase    : Optional[Any]=0.0_2     ,	UpperCAmelCase    : Dict=3     ,	UpperCAmelCase    : int=4     ,	UpperCAmelCase    : List[str]=None     ,	)     ->      Dict:
     lowerCamelCase__       : Optional[Any]   					=						parent
     lowerCamelCase__       : Union[str, Any]   					=						batch_size
     lowerCamelCase__       : int   					=						seq_length
     lowerCamelCase__       : Optional[Any]   					=						is_training
     lowerCamelCase__       : Any   					=						use_input_mask
     lowerCamelCase__       : Dict   					=						use_token_type_ids
     lowerCamelCase__       : Dict   					=						use_labels
     lowerCamelCase__       : Optional[Any]   					=						vocab_size
     lowerCamelCase__       : List[Any]   					=						hidden_size
     lowerCamelCase__       : List[Any]   					=						num_hidden_layers
     lowerCamelCase__       : Optional[int]   					=						num_attention_heads
     lowerCamelCase__       : Tuple   					=						intermediate_size
     lowerCamelCase__       : List[str]   					=						hidden_act
     lowerCamelCase__       : Optional[int]   					=						hidden_dropout_prob
     lowerCamelCase__       : Optional[Any]   					=						attention_probs_dropout_prob
     lowerCamelCase__       : Any   					=						max_position_embeddings
     lowerCamelCase__       : Any   					=						type_vocab_size
     lowerCamelCase__       : Any   					=						type_sequence_label_size
     lowerCamelCase__       : Tuple   					=						initializer_range
     lowerCamelCase__       : Union[str, Any]   					=						num_labels
     lowerCamelCase__       : str   					=						num_choices
     lowerCamelCase__       : int   					=						scope
 def 	A_      ( self    : Dict      )     ->      Any:
     lowerCamelCase__       : Optional[int]   					=						ids_tensor([self.batch_size, self.seq_length]     ,	self.vocab_size      )

     lowerCamelCase__       : str   					=						None
     if self.use_input_mask:
         lowerCamelCase__       : Tuple   					=						random_attention_mask([self.batch_size, self.seq_length]      )

     lowerCamelCase__       : Any   					=						None

     lowerCamelCase__       : Tuple   					=						None
     lowerCamelCase__       : Optional[Any]   					=						None
     lowerCamelCase__       : Optional[int]   					=						None
     if self.use_labels:
         lowerCamelCase__       : List[Any]   					=						ids_tensor([self.batch_size]     ,	self.type_sequence_label_size      )
         lowerCamelCase__       : List[Any]   					=						ids_tensor([self.batch_size, self.seq_length]     ,	self.num_labels      )
         lowerCamelCase__       : Tuple   					=						ids_tensor([self.batch_size]     ,	self.num_choices      )

     lowerCamelCase__       : Optional[int]   					=						self.get_config()

     return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
 def 	A_      ( self    : Any      )     ->      Any:
     return FalconConfig(
         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=UpperCAmelCase     ,	initializer_range=self.initializer_range     ,	pad_token_id=1     ,	new_decoder_architecture=UpperCAmelCase     ,	)
 def 	A_      ( self    : Tuple     ,	UpperCAmelCase    : Optional[Any]     ,	UpperCAmelCase    : Any     ,	UpperCAmelCase    : Dict     ,	UpperCAmelCase    : Optional[int]     ,	UpperCAmelCase    : Tuple     ,	UpperCAmelCase    : List[str]     ,	UpperCAmelCase    : Union[str, Any]      )     ->      Tuple:
     lowerCamelCase__       : Tuple   					=						FalconModel(config=UpperCAmelCase      )
     model.to(UpperCAmelCase      )
     model.eval()
     lowerCamelCase__       : List[str]   					=						model(UpperCAmelCase     ,	attention_mask=UpperCAmelCase      )
     lowerCamelCase__       : List[Any]   					=						model(UpperCAmelCase      )
     self.parent.assertEqual(result.last_hidden_state.shape     ,	(self.batch_size, self.seq_length, self.hidden_size)      )
 def 	A_      ( self    : List[Any]     ,	UpperCAmelCase    : int     ,	UpperCAmelCase    : Any     ,	UpperCAmelCase    : Optional[int]     ,	UpperCAmelCase    : Union[str, Any]     ,	UpperCAmelCase    : Optional[Any]     ,	UpperCAmelCase    : Optional[Any]     ,	UpperCAmelCase    : str     ,	UpperCAmelCase    : Any     ,	UpperCAmelCase    : Union[str, Any]     ,	)     ->      Optional[Any]:
     lowerCamelCase__       : Optional[int]   					=						True
     lowerCamelCase__       : Optional[Any]   					=						FalconModel(UpperCAmelCase      )
     model.to(UpperCAmelCase      )
     model.eval()
     lowerCamelCase__       : str   					=						model(
         UpperCAmelCase     ,	attention_mask=UpperCAmelCase     ,	encoder_hidden_states=UpperCAmelCase     ,	encoder_attention_mask=UpperCAmelCase     ,	)
     lowerCamelCase__       : Union[str, Any]   					=						model(
         UpperCAmelCase     ,	attention_mask=UpperCAmelCase     ,	encoder_hidden_states=UpperCAmelCase     ,	)
     lowerCamelCase__       : Dict   					=						model(UpperCAmelCase     ,	attention_mask=UpperCAmelCase      )
     self.parent.assertEqual(result.last_hidden_state.shape     ,	(self.batch_size, self.seq_length, self.hidden_size)      )
 def 	A_      ( self    : str     ,	UpperCAmelCase    : Optional[Any]     ,	UpperCAmelCase    : Any     ,	UpperCAmelCase    : Any     ,	UpperCAmelCase    : Optional[int]     ,	UpperCAmelCase    : Dict     ,	UpperCAmelCase    : Dict     ,	UpperCAmelCase    : Tuple     ,	UpperCAmelCase    : List[str]     ,	UpperCAmelCase    : str     ,	)     ->      List[str]:
     lowerCamelCase__       : Optional[Any]   					=						FalconForCausalLM(config=UpperCAmelCase      )
     model.to(UpperCAmelCase      )
     model.eval()
     lowerCamelCase__       : int   					=						model(UpperCAmelCase     ,	attention_mask=UpperCAmelCase     ,	labels=UpperCAmelCase      )
     self.parent.assertEqual(result.logits.shape     ,	(self.batch_size, self.seq_length, self.vocab_size)      )
 def 	A_      ( self    : Union[str, Any]     ,	UpperCAmelCase    : Union[str, Any]     ,	UpperCAmelCase    : List[str]     ,	UpperCAmelCase    : Tuple     ,	UpperCAmelCase    : str     ,	UpperCAmelCase    : Optional[int]     ,	UpperCAmelCase    : Optional[int]     ,	UpperCAmelCase    : str     ,	UpperCAmelCase    : Dict     ,	UpperCAmelCase    : Optional[Any]     ,	)     ->      List[str]:
     lowerCamelCase__       : int   					=						True
     lowerCamelCase__       : Optional[int]   					=						True
     lowerCamelCase__       : str   					=						FalconForCausalLM(config=UpperCAmelCase      )
     model.to(UpperCAmelCase      )
     model.eval()

     # first forward pass
     lowerCamelCase__       : List[str]   					=						model(
         UpperCAmelCase     ,	attention_mask=UpperCAmelCase     ,	encoder_hidden_states=UpperCAmelCase     ,	encoder_attention_mask=UpperCAmelCase     ,	use_cache=UpperCAmelCase     ,	)
     lowerCamelCase__       : Dict   					=						outputs.past_key_values

     # create hypothetical multiple next token and extent to next_input_ids
     lowerCamelCase__       : List[str]   					=						ids_tensor((self.batch_size, 3)     ,	config.vocab_size      )
     lowerCamelCase__       : List[Any]   					=						ids_tensor((self.batch_size, 3)     ,	vocab_size=2      )

     # append to next input_ids and
     lowerCamelCase__       : str   					=						torch.cat([input_ids, next_tokens]     ,	dim=-1      )
     lowerCamelCase__       : Union[str, Any]   					=						torch.cat([input_mask, next_mask]     ,	dim=-1      )

     lowerCamelCase__       : Union[str, Any]   					=						model(
         UpperCAmelCase     ,	attention_mask=UpperCAmelCase     ,	encoder_hidden_states=UpperCAmelCase     ,	encoder_attention_mask=UpperCAmelCase     ,	output_hidden_states=UpperCAmelCase     ,	)['hidden_states'][0]
     lowerCamelCase__       : Tuple   					=						model(
         UpperCAmelCase     ,	attention_mask=UpperCAmelCase     ,	encoder_hidden_states=UpperCAmelCase     ,	encoder_attention_mask=UpperCAmelCase     ,	past_key_values=UpperCAmelCase     ,	output_hidden_states=UpperCAmelCase     ,	)['hidden_states'][0]

     # select random slice
     lowerCamelCase__       : Dict   					=						ids_tensor((1,)     ,	output_from_past.shape[-1]      ).item()
     lowerCamelCase__       : List[Any]   					=						output_from_no_past[:, -3:, random_slice_idx].detach()
     lowerCamelCase__       : Dict   					=						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(UpperCAmelCase     ,	UpperCAmelCase     ,	atol=1e-3      )      )







 def 	A_      ( self    : Optional[Any]      )     ->      Optional[int]:
     lowerCamelCase__       : Any   					=						self.prepare_config_and_inputs()
     (
         (
         lowerCamelCase__
     )					,				(
         lowerCamelCase__
     )					,				(
         lowerCamelCase__
     )					,				(
         lowerCamelCase__
     )					,				(
         lowerCamelCase__
     )					,				(
         lowerCamelCase__
     )					,				(
         lowerCamelCase__
     )					,				
     )       : Union[str, Any]   					=						config_and_inputs
     lowerCamelCase__       : List[str]   					=						{'input_ids': input_ids, 'attention_mask': input_mask}
     return config, inputs_dict


@require_torch
class   lowerCAmelCase				(   __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    unittest.TestCase     ):
 UpperCAmelCase__						=						(
     (
         FalconModel,
         FalconForCausalLM,
         FalconForSequenceClassification,
         FalconForTokenClassification,
         FalconForQuestionAnswering,
     )
     if is_torch_available()
     else ()
 )
 UpperCAmelCase__						=						(FalconForCausalLM,) if is_torch_available() else ()
 UpperCAmelCase__						=						(
     {
         """feature-extraction""": FalconModel,
         """text-classification""": FalconForSequenceClassification,
         """text-generation""": FalconForCausalLM,
         """question-answering""": FalconForQuestionAnswering,
         """token-classification""": FalconForTokenClassification,
         """zero-shot""": FalconForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )
 UpperCAmelCase__						=						False
 UpperCAmelCase__						=						False
 def 	A_      ( self    : List[str]      )     ->      List[Any]:
     lowerCamelCase__       : Union[str, Any]   					=						FalconModelTester(self      )
     lowerCamelCase__       : List[Any]   					=						ConfigTester(self     ,	config_class=UpperCAmelCase     ,	hidden_size=37      )
 def 	A_      ( self    : Optional[int]      )     ->      int:
     self.config_tester.run_common_tests()
 def 	A_      ( self    : Dict      )     ->      Dict:
     lowerCamelCase__       : int   					=						self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*UpperCAmelCase      )
 def 	A_      ( self    : List[str]      )     ->      int:
     lowerCamelCase__					,				*lowerCamelCase__       : List[Any]   					=						self.model_tester.prepare_config_and_inputs()
     for alibi in [True, False]:
         lowerCamelCase__       : List[Any]   					=						alibi
         self.model_tester.create_and_check_model(UpperCAmelCase     ,	*UpperCAmelCase      )
 def 	A_      ( self    : Optional[int]      )     ->      Optional[int]:
     lowerCamelCase__					,				lowerCamelCase__       : Any   					=						self.model_tester.prepare_config_and_inputs_for_common()
     lowerCamelCase__       : int   					=						3
     lowerCamelCase__       : List[str]   					=						input_dict['input_ids']
     lowerCamelCase__       : Dict   					=						input_ids.ne(1      ).to(UpperCAmelCase      )
     lowerCamelCase__       : Optional[Any]   					=						ids_tensor([self.model_tester.batch_size]     ,	self.model_tester.type_sequence_label_size      )
     lowerCamelCase__       : int   					=						FalconForSequenceClassification(UpperCAmelCase      )
     model.to(UpperCAmelCase      )
     model.eval()
     lowerCamelCase__       : List[str]   					=						model(UpperCAmelCase     ,	attention_mask=UpperCAmelCase     ,	labels=UpperCAmelCase      )
     self.assertEqual(result.logits.shape     ,	(self.model_tester.batch_size, self.model_tester.num_labels)      )
 def 	A_      ( self    : Optional[Any]      )     ->      Optional[Any]:
     lowerCamelCase__					,				lowerCamelCase__       : List[Any]   					=						self.model_tester.prepare_config_and_inputs_for_common()
     lowerCamelCase__       : Tuple   					=						3
     lowerCamelCase__       : List[str]   					=						'single_label_classification'
     lowerCamelCase__       : Optional[int]   					=						input_dict['input_ids']
     lowerCamelCase__       : Union[str, Any]   					=						input_ids.ne(1      ).to(UpperCAmelCase      )
     lowerCamelCase__       : List[Any]   					=						ids_tensor([self.model_tester.batch_size]     ,	self.model_tester.type_sequence_label_size      )
     lowerCamelCase__       : Union[str, Any]   					=						FalconForSequenceClassification(UpperCAmelCase      )
     model.to(UpperCAmelCase      )
     model.eval()
     lowerCamelCase__       : Any   					=						model(UpperCAmelCase     ,	attention_mask=UpperCAmelCase     ,	labels=UpperCAmelCase      )
     self.assertEqual(result.logits.shape     ,	(self.model_tester.batch_size, self.model_tester.num_labels)      )
 def 	A_      ( self    : Any      )     ->      Union[str, Any]:
     lowerCamelCase__					,				lowerCamelCase__       : List[Any]   					=						self.model_tester.prepare_config_and_inputs_for_common()
     lowerCamelCase__       : List[str]   					=						input_dict['input_ids']
     lowerCamelCase__       : Optional[Any]   					=						FalconForCausalLM(UpperCAmelCase      )
     model.to(UpperCAmelCase      )
     model.eval()
     lowerCamelCase__       : Any   					=						model(UpperCAmelCase     ,	use_cache=UpperCAmelCase      )
     lowerCamelCase__       : int   					=						input_ids.shape[0]
     lowerCamelCase__       : str   					=						model._convert_to_rw_cache(result.past_key_values      )
     lowerCamelCase__       : Optional[Any]   					=						model._convert_cache_to_standard_format(UpperCAmelCase     ,	UpperCAmelCase      )
     for layer in range(len(UpperCAmelCase      )      ):
         for tensor_idx in range(2      ):
             self.assertTrue(rw_cache[layer][tensor_idx].ndim == 3      )
             self.assertTrue(result.past_key_values[layer][tensor_idx].ndim == 4      )
             self.assertTrue(
                 torch.all(result.past_key_values[layer][tensor_idx] == standard_cache[layer][tensor_idx]      )      )
 def 	A_      ( self    : List[str]      )     ->      str:
     lowerCamelCase__					,				lowerCamelCase__       : int   					=						self.model_tester.prepare_config_and_inputs_for_common()
     lowerCamelCase__       : Optional[Any]   					=						3
     lowerCamelCase__       : List[str]   					=						'multi_label_classification'
     lowerCamelCase__       : Optional[int]   					=						input_dict['input_ids']
     lowerCamelCase__       : Tuple   					=						input_ids.ne(1      ).to(UpperCAmelCase      )
     lowerCamelCase__       : int   					=						ids_tensor(
         [self.model_tester.batch_size, config.num_labels]     ,	self.model_tester.type_sequence_label_size      ).to(torch.float      )
     lowerCamelCase__       : Any   					=						FalconForSequenceClassification(UpperCAmelCase      )
     model.to(UpperCAmelCase      )
     model.eval()
     lowerCamelCase__       : Any   					=						model(UpperCAmelCase     ,	attention_mask=UpperCAmelCase     ,	labels=UpperCAmelCase      )
     self.assertEqual(result.logits.shape     ,	(self.model_tester.batch_size, self.model_tester.num_labels)      )







 def 	A_      ( self    : Any      )     ->      Tuple:
     # Falcon can have different numbers of KV-heads than the number of query heads, so we need
     # to override this test to use the right head counts.
     for model_class in self.all_generative_model_classes:
         lowerCamelCase__					,				lowerCamelCase__       : Any   					=						self.model_tester.prepare_config_and_inputs_for_common()

         # If it doesn't support cache, pass the test
         if not hasattr(UpperCAmelCase     ,	'use_cache'      ):
             return

         lowerCamelCase__       : Dict   					=						model_class(UpperCAmelCase      ).to(UpperCAmelCase      )
         if "use_cache" not in inputs:
             lowerCamelCase__       : Optional[Any]   					=						True
         lowerCamelCase__       : Any   					=						model(**UpperCAmelCase      )

         # If "past_key_values" is not returned, pass the test (e.g. RWKV uses a different cache name and format)
         if "past_key_values" not in outputs:
             return

         lowerCamelCase__       : Any   					=						(
             getattr(UpperCAmelCase     ,	'decoder_layers'     ,	UpperCAmelCase      )
             or getattr(UpperCAmelCase     ,	'num_decoder_layers'     ,	UpperCAmelCase      )
             or config.num_hidden_layers
         )
         lowerCamelCase__       : Dict   					=						getattr(UpperCAmelCase     ,	'num_kv_heads'     ,	config.num_attention_heads      )
         lowerCamelCase__       : Dict   					=						getattr(UpperCAmelCase     ,	'd_model'     ,	config.hidden_size      )
         lowerCamelCase__       : str   					=						embed_dim // num_attention_heads

         lowerCamelCase__       : List[Any]   					=						outputs['past_key_values']
         self.assertEqual(len(UpperCAmelCase      )     ,	UpperCAmelCase      )

         lowerCamelCase__					,				lowerCamelCase__       : Any   					=						inputs['input_ids'].shape
         for i in range(UpperCAmelCase      ):
             if config.new_decoder_architecture:
                 lowerCamelCase__       : int   					=						config.num_attention_heads
             elif config.multi_query:
                 lowerCamelCase__       : Optional[Any]   					=						1
             self.assertEqual(len(past_kv[0]      )     ,	2      )  # K V for the decoder = 2
             self.assertEqual(
                 past_kv[i][0].shape     ,	(batch_size, num_attention_heads, seq_length, per_head_embed_dim)      )
             self.assertEqual(
                 past_kv[i][1].shape     ,	(batch_size, num_attention_heads, seq_length, per_head_embed_dim)      )


@require_torch
class   lowerCAmelCase				(   unittest.TestCase     ):
 @slow
 def 	A_      ( self    : int      )     ->      Union[str, Any]:
     lowerCamelCase__       : str   					=						AutoTokenizer.from_pretrained('Rocketknight1/falcon-rw-1b'      )
     lowerCamelCase__       : List[Any]   					=						FalconForCausalLM.from_pretrained('Rocketknight1/falcon-rw-1b'      )
     model.eval()
     model.to(UpperCAmelCase      )
     lowerCamelCase__       : Any   					=						tokenizer('My favorite food is'     ,	return_tensors='pt'      ).to(UpperCAmelCase      )

     lowerCamelCase__       : Optional[Any]   					=						(
         'My favorite food is pizza. I love it so much that I have a pizza party every year for my birthday.'
     )

     lowerCamelCase__       : Optional[Any]   					=						model.generate(**UpperCAmelCase     ,	do_sample=UpperCAmelCase     ,	max_new_tokens=19      )
     lowerCamelCase__       : Union[str, Any]   					=						tokenizer.batch_decode(UpperCAmelCase      )[0]

     self.assertEqual(UpperCAmelCase     ,	UpperCAmelCase      )
 @slow
 def 	A_      ( self    : Any      )     ->      int:
     # The big models are way too big for the CI, so we use tiny random models that resemble their
     # architectures but with much smaller and fewer layers
     for repo in ["Rocketknight1/tiny-random-falcon-7b", "Rocketknight1/tiny-random-falcon-40b"]:
         lowerCamelCase__       : str   					=						AutoTokenizer.from_pretrained(UpperCAmelCase      )
         lowerCamelCase__       : int   					=						FalconForCausalLM.from_pretrained(UpperCAmelCase      )
         model.eval()
         model.to(UpperCAmelCase      )
         lowerCamelCase__       : Tuple   					=						tokenizer('My favorite food is'     ,	return_tensors='pt'      ).to(UpperCAmelCase      )

         # We just test that these run without errors - the models are randomly initialized
         # and so the actual text outputs will be garbage
         model.generate(**UpperCAmelCase     ,	do_sample=UpperCAmelCase     ,	max_new_tokens=4      )
         model.generate(**UpperCAmelCase     ,	do_sample=UpperCAmelCase     ,	max_new_tokens=4      )
         model.generate(**UpperCAmelCase     ,	num_beams=2     ,	max_new_tokens=4      )







 @slow
 def 	A_      ( self    : Union[str, Any]      )     ->      str:
     # The big models are way too big for the CI, so we use tiny random models that resemble their
     # architectures but with much smaller and fewer layers
     with torch.no_grad():
         for repo in [
             "Rocketknight1/falcon-rw-1b",
             "Rocketknight1/tiny-random-falcon-7b",
             "Rocketknight1/tiny-random-falcon-40b",
         ]:
             lowerCamelCase__       : Dict   					=						AutoTokenizer.from_pretrained(UpperCAmelCase      )
             lowerCamelCase__       : Tuple   					=						FalconForCausalLM.from_pretrained(UpperCAmelCase      )
             model.eval()
             model.to(device=UpperCAmelCase      )
             lowerCamelCase__       : Any   					=						tokenizer('My favorite food is'     ,	return_tensors='pt'      ).to(UpperCAmelCase      )

             # Test results are the same with and without cache
             lowerCamelCase__       : List[str]   					=						model.generate(**UpperCAmelCase     ,	do_sample=UpperCAmelCase     ,	max_new_tokens=20     ,	use_cache=UpperCAmelCase      )
             lowerCamelCase__       : Union[str, Any]   					=						model.generate(**UpperCAmelCase     ,	do_sample=UpperCAmelCase     ,	max_new_tokens=20     ,	use_cache=UpperCAmelCase      )
             self.assertTrue((outputs_cache - outputs_no_cache).sum().item() == 0      )





