

'''simple docstring'''




import logging
import os
import sys
from dataclasses import dataclass, field
from typing import List, Optional

import torch
from datasets import Dataset
from torch import nn
from tqdm.auto import tqdm

from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    set_seed,
    utils,
)
from transformers.trainer_utils import get_last_checkpoint, is_main_process


_a :       Optional[Any]       							=  '\nDistills an NLI-based zero-shot classifier to a smaller, more efficient model with a fixed set of candidate class\nnames. Useful for speeding up zero-shot classification in cases where labeled training data is not available, but\nwhen only a single fixed set of classes is needed. Takes a teacher NLI model, student classifier model, unlabeled\ndataset, and set of K possible class names. Yields a single classifier with K outputs corresponding to the provided\nclass names.\n'

_a :       str       							=  logging.getLogger(__name__)



@dataclass
class 				lowercase_		:

  '''simple docstring'''





  __lowerCAmelCase    :       Optional[str]			     =      field(
      default="roberta-large-mnli"       ,				metadata={"help": "The NLI/zero-shot teacher model to be distilled."}       )
  __lowerCAmelCase    :       Optional[str]			     =      field(
      default="This example is {}."       ,				metadata={
          "help": (
              "Template used to turn class names into mock hypotheses for teacher NLI model. Must include {{}}"
              "where class name is inserted."
          )
      }       ,				)
  __lowerCAmelCase    :       Optional[int]			     =      field(
      default=32       ,				metadata={"help": "Batch size for generating teacher predictions."}       )
  __lowerCAmelCase    :       Optional[bool]			     =      field(
      default=_UpperCamelCase       ,				metadata={
          "help": (
              "Allow multiple classes to be true rather than forcing them to sum to 1 (sometimes called"
              "multi-class multi-label classification)."
          )
      }       ,				)
  __lowerCAmelCase    :       Optional[float]			     =      field(
      default=1.0       ,				metadata={"help": "Temperature applied to teacher softmax for distillation."}       )




@dataclass
class 				lowercase_		:

  '''simple docstring'''





  __lowerCAmelCase    :       Optional[str]			     =      field(
      default="distilbert-base-uncased"       ,				metadata={"help": "The NLI/zero-shot teacher model to be distilled."}       )




@dataclass
class 				lowercase_		:

  '''simple docstring'''





  __lowerCAmelCase    :       str			     =      field(metadata={"help": "Text file with one unlabeled instance per line."}       )
  __lowerCAmelCase    :       str			     =      field(metadata={"help": "Text file with one class name per line."}       )
  __lowerCAmelCase    :       bool			     =      field(
      default=_UpperCamelCase       ,				metadata={"help": "Whether to use one of the fast tokenizer (backed by the Rust tokenizers library) or not."}       ,				)




@dataclass
class 				lowercase_		(						_UpperCamelCase       ):

  '''simple docstring'''





  __lowerCAmelCase    :       Optional[str]			     =      field(
      default=_UpperCamelCase       ,				metadata={"help": "The output directory where the model predictions and checkpoints will be written."}       ,				)
  __lowerCAmelCase    :       int			     =      field(
      default=32       ,				metadata={"help": "Batch size per GPU/TPU core/CPU for training."}       )
  __lowerCAmelCase    :       int			     =      field(
      default=1_28       ,				metadata={"help": "Batch size per GPU/TPU core/CPU for evaluation."}       )
  __lowerCAmelCase    :       float			     =      field(default=1.0       ,				metadata={"help": "Total number of training epochs to perform."}       )
  __lowerCAmelCase    :       bool			     =      field(default=_UpperCamelCase       ,				metadata={"help": "Whether to run training of student model."}       )
  __lowerCAmelCase    :       bool			     =      field(
      default=_UpperCamelCase       ,				metadata={
          "help": (
              "Whether to evaluate the agreement of the final student predictions and the teacher predictions"
              "after training."
          )
      }       ,				)
  __lowerCAmelCase    :       Optional[int]			     =      field(
      default=0       ,				metadata={
          "help": (
              "Limit the total amount of checkpoints."
              "Deletes the older checkpoints in the output_dir. Default is 0 (no checkpoints)."
          )
      }       ,				)




class 				lowercase_		(						_UpperCamelCase       ):

  '''simple docstring'''





  def   snake_case_	(   self      ,       a_      ,       a_      ,       a_=False   )							-> Dict:


    """simple docstring"""
    UpperCAmelCase  							=		inputs['labels']
    UpperCAmelCase  							=		model(inputs['input_ids']      ,       attention_mask=inputs['attention_mask']   )
    UpperCAmelCase  							=		outputs[0]

    UpperCAmelCase  							=		-torch.sum(target_p * logits.log_softmax(dim=-1   )      ,       axis=-1   ).mean()

    if return_outputs:
      return loss, outputs

    return loss




def 					lowerCamelCase__						(  SCREAMING_SNAKE_CASE					:      Optional[Any]  ):
  UpperCAmelCase  							=		[]
  with open(_UpperCamelCase				,      'r'  ) as f:
    for line in f:
      UpperCAmelCase  							=		line.strip()
      if len(_UpperCamelCase  ) > 0:
        lines.append(_UpperCamelCase  )
  return lines


def 					lowerCamelCase__						(  SCREAMING_SNAKE_CASE					:      Dict				,      SCREAMING_SNAKE_CASE					:      Dict				,      SCREAMING_SNAKE_CASE					:      Tuple  ):
  UpperCAmelCase  							=		[]
  UpperCAmelCase  							=		[]
  for example in examples:
    for name in class_names:
      premises.append(_UpperCamelCase  )
      hypotheses.append(hypothesis_template.format(_UpperCamelCase  )  )
  return premises, hypotheses


def 					lowerCamelCase__						(  SCREAMING_SNAKE_CASE					:      Tuple  ):
  for label, ind in config.labelaid.items():
    if label.lower().startswith('entail'  ):
      return ind
  logger.warning('Could not identify entailment dimension from teacher config label2id. Setting to -1.'  )
  return -1


def 					lowerCamelCase__						(  SCREAMING_SNAKE_CASE					:      str				,      SCREAMING_SNAKE_CASE					:      List[str]				,      SCREAMING_SNAKE_CASE					:      List[str]				,      SCREAMING_SNAKE_CASE					:      str				,      SCREAMING_SNAKE_CASE					:      int				,      SCREAMING_SNAKE_CASE					:      float				,      SCREAMING_SNAKE_CASE					:      bool				,      SCREAMING_SNAKE_CASE					:      bool				,      SCREAMING_SNAKE_CASE					:      bool				,      SCREAMING_SNAKE_CASE					:      bool				,      ):
  UpperCAmelCase  							=		AutoModelForSequenceClassification.from_pretrained(_UpperCamelCase  )
  UpperCAmelCase  							=		model.config
  if not no_cuda and torch.cuda.is_available():
    UpperCAmelCase  							=		nn.DataParallel(model.cuda()  )
    batch_size *= len(model.device_ids  )
  UpperCAmelCase  							=		AutoTokenizer.from_pretrained(_UpperCamelCase				,      use_fast=_UpperCamelCase  )

  UpperCAmelCase							,					UpperCAmelCase  							=		get_premise_hypothesis_pairs(_UpperCamelCase				,      _UpperCamelCase				,      _UpperCamelCase  )
  UpperCAmelCase  							=		[]

  for i in tqdm(range(0				,      len(_UpperCamelCase  )				,      _UpperCamelCase  )  ):
    UpperCAmelCase  							=		premises[i : i + batch_size]
    UpperCAmelCase  							=		hypotheses[i : i + batch_size]

    UpperCAmelCase  							=		tokenizer(
        _UpperCamelCase				,      _UpperCamelCase				,      padding=_UpperCamelCase				,      truncation='only_first'				,      return_tensors='pt'				,      )

    with torch.cuda.amp.autocast(enabled=_UpperCamelCase  ):
      with torch.no_grad():
        UpperCAmelCase  							=		model(**_UpperCamelCase  )
    logits.append(outputs.logits.detach().cpu().float()  )

  UpperCAmelCase  							=		get_entailment_id(_UpperCamelCase  )
  UpperCAmelCase  							=		-1 if entail_id == 0 else 0
  UpperCAmelCase  							=		torch.cat(_UpperCamelCase				,      dim=0  )  # N*K x 3
  UpperCAmelCase  							=		logits.reshape(len(_UpperCamelCase  )				,      len(_UpperCamelCase  )				,      -1  )[..., [contr_id, entail_id]]  # N x K x 2

  if multi_label:
    # softmax over (contr, entail) logits for each class independently
    UpperCAmelCase  							=		(nli_logits / temperature).softmax(-1  )
  else:
    # softmax over entail logits across classes s.t. class probabilities sum to 1.
    UpperCAmelCase  							=		(nli_logits / temperature).softmax(1  )

  return nli_prob[..., 1]  # N x K


def 					lowerCamelCase__						(  ):
  UpperCAmelCase  							=		HfArgumentParser(
      (DataTrainingArguments, TeacherModelArguments, StudentModelArguments, DistillTrainingArguments)				,      description=_UpperCamelCase				,      )

  if len(sys.argv  ) == 2 and sys.argv[1].endswith('.json'  ):
    # If we pass only one argument to the script and it's the path to a json file,
    # let's parse it to get our arguments.
    UpperCAmelCase							,					UpperCAmelCase							,					UpperCAmelCase							,					UpperCAmelCase  							=		parser.parse_json_file(
        json_file=os.path.abspath(sys.argv[1]  )  )
  else:
    UpperCAmelCase							,					UpperCAmelCase							,					UpperCAmelCase							,					UpperCAmelCase  							=		parser.parse_args_into_dataclasses()

  # Detecting last checkpoint.
  UpperCAmelCase  							=		None
  if os.path.isdir(training_args.output_dir  ) and training_args.do_train and not training_args.overwrite_output_dir:
    UpperCAmelCase  							=		get_last_checkpoint(training_args.output_dir  )
    if last_checkpoint is None and len(os.listdir(training_args.output_dir  )  ) > 0:
      raise ValueError(
          f'''Output directory ({training_args.output_dir}) already exists and is not empty. '''
          'Use --overwrite_output_dir to overcome.'  )
    elif last_checkpoint is not None:
      logger.info(
          f'''Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change '''
          'the `--output_dir` or add `--overwrite_output_dir` to train from scratch.'  )

    # Setup logging
  logging.basicConfig(
      format='%(asctime)s - %(levelname)s - %(name)s - %(message)s'				,      datefmt='%m/%d/%Y %H:%M:%S'				,      handlers=[logging.StreamHandler(sys.stdout  )]				,      )
  logger.setLevel(logging.INFO if is_main_process(training_args.local_rank  ) else logging.WARN  )

  # Log on each process the small summary:
  logger.warning(
      f'''Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}'''
      + f'''distributed training: {bool(training_args.local_rank != -1  )}, 16-bits training: {training_args.fpaa}'''  )
  # Set the verbosity to info of the Transformers logger (on main process only):
  if is_main_process(training_args.local_rank  ):
    utils.logging.set_verbosity_info()
    utils.logging.enable_default_handler()
    utils.logging.enable_explicit_format()

  if training_args.local_rank != -1:
    raise ValueError('Distributed training is not currently supported.'  )
  if training_args.tpu_num_cores is not None:
    raise ValueError('TPU acceleration is not currently supported.'  )

  logger.info(f'''Training/evaluation parameters {training_args}'''  )

  # Set seed before initializing model.
  set_seed(training_args.seed  )

  # 1. read in data
  UpperCAmelCase  							=		read_lines(data_args.data_file  )
  UpperCAmelCase  							=		read_lines(data_args.class_names_file  )

  # 2. get teacher predictions and load into dataset
  logger.info('Generating predictions from zero-shot teacher model'  )
  UpperCAmelCase  							=		get_teacher_predictions(
      teacher_args.teacher_name_or_path				,      _UpperCamelCase				,      _UpperCamelCase				,      teacher_args.hypothesis_template				,      teacher_args.teacher_batch_size				,      teacher_args.temperature				,      teacher_args.multi_label				,      data_args.use_fast_tokenizer				,      training_args.no_cuda				,      training_args.fpaa				,      )
  UpperCAmelCase  							=		Dataset.from_dict(
      {
          'text': examples,
          'labels': teacher_soft_preds,
      }  )

  # 3. create student
  logger.info('Initializing student model'  )
  UpperCAmelCase  							=		AutoModelForSequenceClassification.from_pretrained(
      student_args.student_name_or_path				,      num_labels=len(_UpperCamelCase  )  )
  UpperCAmelCase  							=		AutoTokenizer.from_pretrained(student_args.student_name_or_path				,      use_fast=data_args.use_fast_tokenizer  )
  UpperCAmelCase  							=		dict(enumerate(_UpperCamelCase  )  )
  UpperCAmelCase  							=		{label: i for i, label in enumerate(_UpperCamelCase  )}

  # 4. train student on teacher predictions
  UpperCAmelCase  							=		dataset.map(_UpperCamelCase				,      input_columns='text'  )
  dataset.set_format('torch'  )

  def compute_metrics(SCREAMING_SNAKE_CASE					:      str				,      SCREAMING_SNAKE_CASE					:      Dict=False  ):
    UpperCAmelCase  							=		p.predictions.argmax(-1  )
    UpperCAmelCase  							=		p.label_ids.argmax(-1  )  # "label_ids" are actually distributions
    return {"agreement": (preds == proxy_labels).mean().item()}

  UpperCAmelCase  							=		DistillationTrainer(
      model=_UpperCamelCase				,      tokenizer=_UpperCamelCase				,      args=_UpperCamelCase				,      train_dataset=_UpperCamelCase				,      compute_metrics=_UpperCamelCase				,      )

  if training_args.do_train:
    logger.info('Training student model on teacher predictions'  )
    trainer.train()

  if training_args.do_eval:
    UpperCAmelCase  							=		trainer.evaluate(eval_dataset=_UpperCamelCase  )['eval_agreement']
    logger.info(f'''Agreement of student and teacher predictions: {agreement * 100:0.2f}%'''  )

  trainer.save_model()


if __name__ == "__main__":
  main()



'''simple docstring'''




import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaTokenizer, GPTNeoConfig, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
  import jax
  import jax.numpy as jnp

  from transformers.modeling_flax_pytorch_utils import (
      convert_pytorch_state_dict_to_flax,
      load_flax_weights_in_pytorch_model,
  )
  from transformers.models.gpt_neo.modeling_flax_gpt_neo import FlaxGPTNeoForCausalLM, FlaxGPTNeoModel

if is_torch_available():
  import torch




class 	__UpperCAmelCase				:





       '''simple docstring'''



       def __init__(      self    ,     _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE=14    ,     _SCREAMING_SNAKE_CASE=7    ,     _SCREAMING_SNAKE_CASE=True    ,     _SCREAMING_SNAKE_CASE=True    ,     _SCREAMING_SNAKE_CASE=False    ,     _SCREAMING_SNAKE_CASE=True    ,     _SCREAMING_SNAKE_CASE=99    ,     _SCREAMING_SNAKE_CASE=32    ,     _SCREAMING_SNAKE_CASE=4    ,     _SCREAMING_SNAKE_CASE=4    ,     _SCREAMING_SNAKE_CASE=[[["global", "local"], 2]]    ,     _SCREAMING_SNAKE_CASE=37    ,     _SCREAMING_SNAKE_CASE="gelu"    ,     _SCREAMING_SNAKE_CASE=0.1    ,     _SCREAMING_SNAKE_CASE=0.1    ,     _SCREAMING_SNAKE_CASE=512    ,     _SCREAMING_SNAKE_CASE=7    ,     _SCREAMING_SNAKE_CASE=0.02    ,     )		->     Dict:
           A_     =   parent
           A_     =   batch_size
           A_     =   seq_length
           A_     =   is_training
           A_     =   use_input_mask
           A_     =   use_token_type_ids
           A_     =   use_labels
           A_     =   vocab_size
           A_     =   hidden_size
           A_     =   num_hidden_layers
           A_     =   num_attention_heads
           A_     =   attention_types
           A_     =   intermediate_size
           A_     =   hidden_act
           A_     =   hidden_dropout_prob
           A_     =   attention_probs_dropout_prob
           A_     =   max_position_embeddings
           A_     =   window_size
           A_     =   initializer_range
           A_     =   None
           A_     =   vocab_size - 1
           A_     =   vocab_size - 1
           A_     =   vocab_size - 1

       def 					__A       (      self )		->     Optional[int]:
           A_     =   ids_tensor([self.batch_size, self.seq_length]    ,     self.vocab_size )

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

           A_     =   GPTNeoConfig(
               vocab_size=self.vocab_size    ,     hidden_size=self.hidden_size    ,     num_layers=self.num_hidden_layers    ,     num_heads=self.num_attention_heads    ,     max_position_embeddings=self.max_position_embeddings    ,     use_cache=_SCREAMING_SNAKE_CASE    ,     bos_token_id=self.bos_token_id    ,     eos_token_id=self.eos_token_id    ,     pad_token_id=self.pad_token_id    ,     window_size=self.window_size    ,     attention_types=self.attention_types    ,     )

           return (config, input_ids, input_mask)

       def 					__A       (      self )		->     str:
           A_     =   self.prepare_config_and_inputs()
           A_   ,A_   ,A_     =   config_and_inputs
           A_     =   {'''input_ids''': input_ids, '''attention_mask''': attention_mask}
           return config, inputs_dict

       def 					__A       (      self    ,     _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE )		->     Union[str, Any]:
           A_     =   20
           A_     =   model_class_name(_SCREAMING_SNAKE_CASE )

           A_     =   model.init_cache(input_ids.shape[0]    ,     _SCREAMING_SNAKE_CASE )
           A_     =   jnp.ones((input_ids.shape[0], max_decoder_length)    ,     dtype='''i4''' )

           A_     =   jnp.broadcast_to(
               jnp.arange(input_ids.shape[-1] - 1 )[None, :]    ,     (input_ids.shape[0], input_ids.shape[-1] - 1) )
           A_     =   model(
               input_ids[:, :-1]    ,     attention_mask=_SCREAMING_SNAKE_CASE    ,     past_key_values=_SCREAMING_SNAKE_CASE    ,     position_ids=_SCREAMING_SNAKE_CASE    ,     )

           A_     =   jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]    ,     dtype='''i4''' )
           A_     =   model(
               input_ids[:, -1:]    ,     attention_mask=_SCREAMING_SNAKE_CASE    ,     past_key_values=outputs_cache.past_key_values    ,     position_ids=_SCREAMING_SNAKE_CASE    ,     )

           A_     =   model(_SCREAMING_SNAKE_CASE )

           A_     =   np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5]) ) )
           self.parent.assertTrue(diff < 1E-3    ,     msg=F'''Max diff is {diff}''' )

       def 					__A       (      self    ,     _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE )		->     List[Any]:
           A_     =   20
           A_     =   model_class_name(_SCREAMING_SNAKE_CASE )

           A_     =   jnp.concatenate(
               [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1]) )]    ,     axis=-1    ,     )

           A_     =   model.init_cache(input_ids.shape[0]    ,     _SCREAMING_SNAKE_CASE )
           A_     =   jnp.broadcast_to(
               jnp.arange(input_ids.shape[-1] - 1 )[None, :]    ,     (input_ids.shape[0], input_ids.shape[-1] - 1) )

           A_     =   model(
               input_ids[:, :-1]    ,     attention_mask=_SCREAMING_SNAKE_CASE    ,     past_key_values=_SCREAMING_SNAKE_CASE    ,     position_ids=_SCREAMING_SNAKE_CASE    ,     )
           A_     =   jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]    ,     dtype='''i4''' )
           A_     =   model(
               input_ids[:, -1:]    ,     past_key_values=outputs_cache.past_key_values    ,     attention_mask=_SCREAMING_SNAKE_CASE    ,     position_ids=_SCREAMING_SNAKE_CASE    ,     )

           A_     =   model(_SCREAMING_SNAKE_CASE    ,     attention_mask=_SCREAMING_SNAKE_CASE )

           A_     =   np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5]) ) )
           self.parent.assertTrue(diff < 1E-3    ,     msg=F'''Max diff is {diff}''' )




@require_flax
class 	__UpperCAmelCase				(				_UpperCamelCase						,      _UpperCamelCase						,      unittest.TestCase     ):





       '''simple docstring'''



       __lowercase					:					List[str]    =				(FlaxGPTNeoModel, FlaxGPTNeoForCausalLM) if is_flax_available() else ()
       __lowercase					:					Dict    =				(FlaxGPTNeoForCausalLM,) if is_flax_available() else ()

       def 					__A       (      self )		->     Any:
           A_     =   FlaxGPTNeoModelTester(self )

       def 					__A       (      self )		->     Optional[int]:
           for model_class_name in self.all_model_classes:
               A_   ,A_   ,A_     =   self.model_tester.prepare_config_and_inputs()
               self.model_tester.check_use_cache_forward(_SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE )

       def 					__A       (      self )		->     Union[str, Any]:
           for model_class_name in self.all_model_classes:
               A_   ,A_   ,A_     =   self.model_tester.prepare_config_and_inputs()
               self.model_tester.check_use_cache_forward_with_attn_mask(
                   _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE )

       @slow
       def 					__A       (      self )		->     Optional[Any]:
           A_     =   GPTaTokenizer.from_pretrained('''gpt2'''    ,     pad_token='''<|endoftext|>'''    ,     padding_side='''left''' )
           A_     =   tokenizer(['''Hello this is a long string''', '''Hey''']    ,     return_tensors='''np'''    ,     padding=_SCREAMING_SNAKE_CASE    ,     truncation=_SCREAMING_SNAKE_CASE )

           A_     =   FlaxGPTNeoForCausalLM.from_pretrained('''EleutherAI/gpt-neo-125M''' )
           A_     =   False
           A_     =   model.config.eos_token_id

           A_     =   jax.jit(model.generate )

           A_     =   jit_generate(
               inputs['''input_ids''']    ,     attention_mask=inputs['''attention_mask''']    ,     pad_token_id=tokenizer.pad_token_id ).sequences

           A_     =   tokenizer.batch_decode(_SCREAMING_SNAKE_CASE    ,     skip_special_tokens=_SCREAMING_SNAKE_CASE )

           A_     =   [
               '''Hello this is a long string of text.\n\nI\'m trying to get the text of the''',
               '''Hey, I\'m a little late to the party. I\'m going to''',
           ]

           self.assertListEqual(_SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE )

       @is_pt_flax_cross_test
       def 					__A       (      self )		->     Union[str, Any]:
           A_   ,A_     =   self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
               with self.subTest(model_class.__name__ ):
                   # prepare inputs
                   A_     =   self._prepare_for_class(_SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE )
                   A_     =   {k: torch.tensor(v.tolist() ) for k, v in prepared_inputs_dict.items()}

                   # load corresponding PyTorch class
                   A_     =   model_class.__name__[4:]  # Skip the "Flax" at the beginning
                   A_     =   getattr(_SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE )

                   A_   ,A_     =   pt_inputs['''input_ids'''].shape
                   A_     =   np.random.randint(0    ,     seq_length - 1    ,     size=(batch_size,) )
                   for batch_idx, start_index in enumerate(_SCREAMING_SNAKE_CASE ):
                       A_     =   0
                       A_     =   1
                       A_     =   0
                       A_     =   1
                   A_     =   pt_model_class(_SCREAMING_SNAKE_CASE ).eval()
                   A_     =   model_class(_SCREAMING_SNAKE_CASE    ,     dtype=jnp.floataa )

                   A_     =   convert_pytorch_state_dict_to_flax(pt_model.state_dict()    ,     _SCREAMING_SNAKE_CASE )
                   A_     =   fx_state

                   with torch.no_grad():
                       A_     =   pt_model(**_SCREAMING_SNAKE_CASE ).to_tuple()

                   A_     =   fx_model(**_SCREAMING_SNAKE_CASE ).to_tuple()
                   self.assertEqual(len(_SCREAMING_SNAKE_CASE )    ,     len(_SCREAMING_SNAKE_CASE )    ,     '''Output lengths differ between Flax and PyTorch''' )
                   for fx_output, pt_output in zip(_SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE ):
                       self.assert_almost_equals(fx_output[:, -1]    ,     pt_output[:, -1].numpy()    ,     4E-2 )

                   with tempfile.TemporaryDirectory() as tmpdirname:
                       pt_model.save_pretrained(_SCREAMING_SNAKE_CASE )
                       A_     =   model_class.from_pretrained(_SCREAMING_SNAKE_CASE    ,     from_pt=_SCREAMING_SNAKE_CASE )

                   A_     =   fx_model_loaded(**_SCREAMING_SNAKE_CASE ).to_tuple()
                   self.assertEqual(
                       len(_SCREAMING_SNAKE_CASE )    ,     len(_SCREAMING_SNAKE_CASE )    ,     '''Output lengths differ between Flax and PyTorch''' )
                   for fx_output_loaded, pt_output in zip(_SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE ):
                       self.assert_almost_equals(fx_output_loaded[:, -1]    ,     pt_output[:, -1].numpy()    ,     4E-2 )

       @is_pt_flax_cross_test
       def 					__A       (      self )		->     Union[str, Any]:
           A_   ,A_     =   self.model_tester.prepare_config_and_inputs_for_common()
           for model_class in self.all_model_classes:
               with self.subTest(model_class.__name__ ):
                   # prepare inputs
                   A_     =   self._prepare_for_class(_SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE )
                   A_     =   {k: torch.tensor(v.tolist() ) for k, v in prepared_inputs_dict.items()}

                   # load corresponding PyTorch class
                   A_     =   model_class.__name__[4:]  # Skip the "Flax" at the beginning
                   A_     =   getattr(_SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE )

                   A_     =   pt_model_class(_SCREAMING_SNAKE_CASE ).eval()
                   A_     =   model_class(_SCREAMING_SNAKE_CASE    ,     dtype=jnp.floataa )

                   A_     =   load_flax_weights_in_pytorch_model(_SCREAMING_SNAKE_CASE    ,     fx_model.params )
                   A_   ,A_     =   pt_inputs['''input_ids'''].shape
                   A_     =   np.random.randint(0    ,     seq_length - 1    ,     size=(batch_size,) )
                   for batch_idx, start_index in enumerate(_SCREAMING_SNAKE_CASE ):
                       A_     =   0
                       A_     =   1
                       A_     =   0
                       A_     =   1

                   # make sure weights are tied in PyTorch
                   pt_model.tie_weights()

                   with torch.no_grad():
                       A_     =   pt_model(**_SCREAMING_SNAKE_CASE ).to_tuple()

                   A_     =   fx_model(**_SCREAMING_SNAKE_CASE ).to_tuple()
                   self.assertEqual(len(_SCREAMING_SNAKE_CASE )    ,     len(_SCREAMING_SNAKE_CASE )    ,     '''Output lengths differ between Flax and PyTorch''' )
                   for fx_output, pt_output in zip(_SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE ):
                       self.assert_almost_equals(fx_output[:, -1]    ,     pt_output[:, -1].numpy()    ,     4E-2 )

                   with tempfile.TemporaryDirectory() as tmpdirname:
                       fx_model.save_pretrained(_SCREAMING_SNAKE_CASE )
                       A_     =   pt_model_class.from_pretrained(_SCREAMING_SNAKE_CASE    ,     from_flax=_SCREAMING_SNAKE_CASE )

                   with torch.no_grad():
                       A_     =   pt_model_loaded(**_SCREAMING_SNAKE_CASE ).to_tuple()

                   self.assertEqual(
                       len(_SCREAMING_SNAKE_CASE )    ,     len(_SCREAMING_SNAKE_CASE )    ,     '''Output lengths differ between Flax and PyTorch''' )
                   for fx_output, pt_output in zip(_SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE ):
                       self.assert_almost_equals(fx_output[:, -1]    ,     pt_output[:, -1].numpy()    ,     4E-2 )

       @slow
       def 					__A       (      self )		->     List[str]:
           for model_class_name in self.all_model_classes:
               A_     =   model_class_name.from_pretrained('''EleutherAI/gpt-neo-125M''' )
               A_     =   model(np.ones((1, 1) ) )
               self.assertIsNotNone(_SCREAMING_SNAKE_CASE )






