





from __future__ import annotations

import unittest

from transformers import AutoTokenizer, GPTJConfig, is_tf_available
from transformers.testing_utils import require_tf, slow, tooslow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin
from ...utils.test_modeling_tf_core import TFCoreModelTesterMixin


if is_tf_available():
				import tensorflow as tf

				from transformers.models.gptj.modeling_tf_gptj import (
				    TFGPTJForCausalLM,
				    TFGPTJForQuestionAnswering,
				    TFGPTJForSequenceClassification,
				    TFGPTJModel,
				    shape_list,
				)




class 	_UpperCAmelCase						:




					def __init__(    self     ,   __lowerCAmelCase							)    ->	int:
							UpperCamelCase__								=						parent
							UpperCamelCase__								=						13
							UpperCamelCase__								=						7
							UpperCamelCase__								=						True
							UpperCamelCase__								=						True
							UpperCamelCase__								=						True
							UpperCamelCase__								=						True
							UpperCamelCase__								=						True
							UpperCamelCase__								=						99
							UpperCamelCase__								=						32
							UpperCamelCase__								=						4
							UpperCamelCase__								=						2
							UpperCamelCase__								=						4
							UpperCamelCase__								=						37
							UpperCamelCase__								=						"""gelu"""
							UpperCamelCase__								=						0.1
							UpperCamelCase__								=						0.1
							UpperCamelCase__								=						512
							UpperCamelCase__								=						16
							UpperCamelCase__								=						2
							UpperCamelCase__								=						0.0_2
							UpperCamelCase__								=						3
							UpperCamelCase__								=						4
							UpperCamelCase__								=						None
							UpperCamelCase__								=						self.vocab_size - 1
							UpperCamelCase__								=						self.vocab_size - 1
							UpperCamelCase__								=						self.vocab_size - 1




					def 						snake_case_		(    self							)    ->	Tuple:
							UpperCamelCase__								=						ids_tensor([self.batch_size, self.seq_length]     ,   self.vocab_size							)

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

							UpperCamelCase__								=						None
							if self.use_token_type_ids:
									UpperCamelCase__								=						ids_tensor([self.batch_size, self.seq_length]     ,   self.type_vocab_size							)

							UpperCamelCase__								=						None
							if self.use_mc_token_ids:
									UpperCamelCase__								=						ids_tensor([self.batch_size, self.num_choices]     ,   self.seq_length							)

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

							UpperCamelCase__								=						GPTJConfig(
							    vocab_size=self.vocab_size     ,   n_embd=self.hidden_size     ,   n_layer=self.num_hidden_layers     ,   n_head=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     ,   n_positions=self.max_position_embeddings     ,   type_vocab_size=self.type_vocab_size     ,   initializer_range=self.initializer_range     ,   bos_token_id=self.bos_token_id     ,   eos_token_id=self.eos_token_id     ,   pad_token_id=self.pad_token_id     ,   rotary_dim=self.rotary_dim     ,   return_dict=__lowerCAmelCase     ,   )

							UpperCamelCase__								=						ids_tensor([self.num_hidden_layers, self.num_attention_heads]     ,   2							)

							return (
							    config,
							    input_ids,
							    input_mask,
							    head_mask,
							    token_type_ids,
							    mc_token_ids,
							    sequence_labels,
							    token_labels,
							    choice_labels,
							)




					def 						snake_case_		(    self     ,   __lowerCAmelCase     ,   __lowerCAmelCase     ,   __lowerCAmelCase     ,   __lowerCAmelCase     ,   __lowerCAmelCase     ,   *__lowerCAmelCase							)    ->	int:
							UpperCamelCase__								=						TFGPTJModel(config=__lowerCAmelCase							)
							UpperCamelCase__								=						{
							    """input_ids""": input_ids,
							    """attention_mask""": input_mask,
							    """token_type_ids""": token_type_ids,
							}
							UpperCamelCase__								=						model(__lowerCAmelCase							)

							UpperCamelCase__								=						[input_ids, None, input_mask]  # None is the input for 'past'
							UpperCamelCase__								=						model(__lowerCAmelCase							)

							UpperCamelCase__								=						model(__lowerCAmelCase							)

							self.parent.assertEqual(result.last_hidden_state.shape     ,   (self.batch_size, self.seq_length, self.hidden_size)							)




					def 						snake_case_		(    self     ,   __lowerCAmelCase     ,   __lowerCAmelCase     ,   __lowerCAmelCase     ,   __lowerCAmelCase     ,   __lowerCAmelCase     ,   *__lowerCAmelCase							)    ->	Optional[int]:
							UpperCamelCase__								=						TFGPTJModel(config=__lowerCAmelCase							)

							# first forward pass
							UpperCamelCase__								=						model(__lowerCAmelCase     ,   token_type_ids=__lowerCAmelCase     ,   use_cache=__lowerCAmelCase							)
							UpperCamelCase__								=						model(__lowerCAmelCase     ,   token_type_ids=__lowerCAmelCase							)
							UpperCamelCase__								=						model(__lowerCAmelCase     ,   token_type_ids=__lowerCAmelCase     ,   use_cache=__lowerCAmelCase							)

							self.parent.assertTrue(len(__lowerCAmelCase							) == len(__lowerCAmelCase							)							)
							self.parent.assertTrue(len(__lowerCAmelCase							) == len(__lowerCAmelCase							) + 1							)

							UpperCamelCase__	,      UpperCamelCase__								=						outputs.to_tuple()

							# create hypothetical next token and extent to next_input_ids
							UpperCamelCase__								=						ids_tensor((self.batch_size, 1)     ,   config.vocab_size							)
							UpperCamelCase__								=						ids_tensor([self.batch_size, 1]     ,   self.type_vocab_size							)

							# append to next input_ids and token_type_ids
							UpperCamelCase__								=						tf.concat([input_ids, next_tokens]     ,   axis=-1							)
							UpperCamelCase__								=						tf.concat([token_type_ids, next_token_types]     ,   axis=-1							)

							UpperCamelCase__								=						model(__lowerCAmelCase     ,   token_type_ids=__lowerCAmelCase							)["""last_hidden_state"""]
							UpperCamelCase__								=						model(__lowerCAmelCase     ,   token_type_ids=__lowerCAmelCase     ,   past_key_values=__lowerCAmelCase							)[
							    """last_hidden_state"""
							]

							# select random slice
							UpperCamelCase__								=						int(ids_tensor((1,)     ,   shape_list(__lowerCAmelCase							)[-1]							)							)
							UpperCamelCase__								=						output_from_no_past[:, -1, random_slice_idx]
							UpperCamelCase__								=						output_from_past[:, 0, random_slice_idx]

							# test that outputs are equal for slice
							tf.debugging.assert_near(__lowerCAmelCase     ,   __lowerCAmelCase     ,   rtol=1E-6							)




					def 						snake_case_		(    self     ,   __lowerCAmelCase     ,   __lowerCAmelCase     ,   __lowerCAmelCase     ,   __lowerCAmelCase     ,   __lowerCAmelCase     ,   *__lowerCAmelCase							)    ->	int:
							UpperCamelCase__								=						TFGPTJModel(config=__lowerCAmelCase							)

							# create attention mask
							UpperCamelCase__								=						self.seq_length // 2
							UpperCamelCase__								=						tf.ones((self.batch_size, half_seq_length)     ,   dtype=tf.intaa							)
							UpperCamelCase__								=						tf.zeros((self.batch_size, self.seq_length - half_seq_length)     ,   dtype=tf.intaa							)
							UpperCamelCase__								=						tf.concat([attn_mask_begin, attn_mask_end]     ,   axis=1							)

							# first forward pass
							UpperCamelCase__	,      UpperCamelCase__								=						model(__lowerCAmelCase     ,   attention_mask=__lowerCAmelCase							).to_tuple()

							# create hypothetical next token and extent to next_input_ids
							UpperCamelCase__								=						ids_tensor((self.batch_size, 1)     ,   config.vocab_size							)

							# change a random masked slice from input_ids
							UpperCamelCase__								=						ids_tensor((1,)     ,   __lowerCAmelCase							).numpy() + 1
							UpperCamelCase__								=						ids_tensor((self.batch_size, self.seq_length)     ,   config.vocab_size							)
							UpperCamelCase__								=						tf.range(self.seq_length							) == (self.seq_length - random_seq_idx_to_change)
							UpperCamelCase__								=						tf.transpose(
							    tf.broadcast_to(tf.expand_dims(__lowerCAmelCase     ,   -1							)     ,   (self.seq_length, self.batch_size)							)							)
							UpperCamelCase__								=						tf.where(__lowerCAmelCase     ,   __lowerCAmelCase     ,   __lowerCAmelCase							)

							# append to next input_ids and attn_mask
							UpperCamelCase__								=						tf.concat([input_ids, next_tokens]     ,   axis=-1							)
							UpperCamelCase__								=						tf.concat([attn_mask, tf.ones((shape_list(__lowerCAmelCase							)[0], 1)     ,   dtype=tf.intaa							)]     ,   axis=1							)

							# get two different outputs
							UpperCamelCase__								=						model(__lowerCAmelCase     ,   attention_mask=__lowerCAmelCase							)["""last_hidden_state"""]
							UpperCamelCase__								=						model(__lowerCAmelCase     ,   past_key_values=__lowerCAmelCase     ,   attention_mask=__lowerCAmelCase							)[
							    """last_hidden_state"""
							]

							# select random slice
							UpperCamelCase__								=						int(ids_tensor((1,)     ,   shape_list(__lowerCAmelCase							)[-1]							)							)
							UpperCamelCase__								=						output_from_no_past[:, -1, random_slice_idx]
							UpperCamelCase__								=						output_from_past[:, 0, random_slice_idx]

							# test that outputs are equal for slice
							tf.debugging.assert_near(__lowerCAmelCase     ,   __lowerCAmelCase     ,   rtol=1E-1_2							)




					def 						snake_case_		(    self     ,   __lowerCAmelCase     ,   __lowerCAmelCase     ,   __lowerCAmelCase     ,   __lowerCAmelCase     ,   __lowerCAmelCase     ,   *__lowerCAmelCase							)    ->	Union[str, Any]:
							UpperCamelCase__								=						TFGPTJModel(config=__lowerCAmelCase							)

							UpperCamelCase__								=						input_ids[:1, :]
							UpperCamelCase__								=						input_mask[:1, :]
							UpperCamelCase__								=						token_type_ids[:1, :]
							UpperCamelCase__								=						1

							# first forward pass
							UpperCamelCase__								=						model(__lowerCAmelCase     ,   attention_mask=__lowerCAmelCase     ,   token_type_ids=__lowerCAmelCase     ,   use_cache=__lowerCAmelCase							)

							UpperCamelCase__	,      UpperCamelCase__								=						outputs.to_tuple()

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

							# append to next input_ids and token_type_ids
							UpperCamelCase__								=						tf.concat([input_ids, next_tokens]     ,   axis=-1							)
							UpperCamelCase__								=						tf.concat([input_mask, next_attn_mask]     ,   axis=-1							)
							UpperCamelCase__								=						tf.concat([token_type_ids, next_token_types]     ,   axis=-1							)

							UpperCamelCase__								=						model(
							    __lowerCAmelCase     ,   token_type_ids=__lowerCAmelCase     ,   attention_mask=__lowerCAmelCase							)["""last_hidden_state"""]
							UpperCamelCase__								=						model(
							    __lowerCAmelCase     ,   token_type_ids=__lowerCAmelCase     ,   attention_mask=__lowerCAmelCase     ,   past_key_values=__lowerCAmelCase     ,   )["""last_hidden_state"""]
							self.parent.assertTrue(output_from_past.shape[1] == next_tokens.shape[1]							)

							# select random slice
							UpperCamelCase__								=						int(ids_tensor((1,)     ,   shape_list(__lowerCAmelCase							)[-1]							)							)
							UpperCamelCase__								=						output_from_no_past[:, -3:, random_slice_idx]
							UpperCamelCase__								=						output_from_past[:, :, random_slice_idx]

							# test that outputs are equal for slice
							tf.debugging.assert_near(__lowerCAmelCase     ,   __lowerCAmelCase     ,   rtol=1E-3							)




					def 						snake_case_		(    self     ,   __lowerCAmelCase     ,   __lowerCAmelCase     ,   __lowerCAmelCase     ,   __lowerCAmelCase     ,   __lowerCAmelCase     ,   *__lowerCAmelCase							)    ->	Dict:
							UpperCamelCase__								=						TFGPTJForCausalLM(config=__lowerCAmelCase							)
							UpperCamelCase__								=						{
							    """input_ids""": input_ids,
							    """attention_mask""": input_mask,
							    """token_type_ids""": token_type_ids,
							}
							UpperCamelCase__								=						model(__lowerCAmelCase							)
							self.parent.assertEqual(result.logits.shape     ,   (self.batch_size, self.seq_length, self.vocab_size)							)




					def 						snake_case_		(    self							)    ->	List[str]:
							UpperCamelCase__								=						self.prepare_config_and_inputs()

							(
							    (
							    UpperCamelCase__
							)	,      (
							    UpperCamelCase__
							)	,      (
							    UpperCamelCase__
							)	,      (
							    UpperCamelCase__
							)	,      (
							    UpperCamelCase__
							)	,      (
							    UpperCamelCase__
							)	,      (
							    UpperCamelCase__
							)	,      (
							    UpperCamelCase__
							)	,      (
							    UpperCamelCase__
							)	,      
							)								=						config_and_inputs

							UpperCamelCase__								=						{
							    """input_ids""": input_ids,
							    """token_type_ids""": token_type_ids,
							    """attention_mask""": input_mask,
							}
							return config, inputs_dict




@require_tf
class 	_UpperCAmelCase						( _lowercase			,							_lowercase			,							_lowercase			,							unittest.TestCase       ):
					__A        = (
					    (TFGPTJForCausalLM, TFGPTJForSequenceClassification, TFGPTJForQuestionAnswering, TFGPTJModel)
					    if is_tf_available()
					    else ()
					)

					__A        = (TFGPTJForCausalLM,) if is_tf_available() else ()
					__A        = (
					    {
					        "feature-extraction": TFGPTJModel,
					        "question-answering": TFGPTJForQuestionAnswering,
					        "text-classification": TFGPTJForSequenceClassification,
					        "text-generation": TFGPTJForCausalLM,
					        "zero-shot": TFGPTJForSequenceClassification,
					    }
					    if is_tf_available()
					    else {}
					)
					__A        = False
					__A        = False
					__A        = False
					__A        = False




					def 						snake_case_		(    self     ,   __lowerCAmelCase     ,   __lowerCAmelCase     ,   __lowerCAmelCase     ,   __lowerCAmelCase     ,   __lowerCAmelCase							)    ->	Optional[Any]:
							if (
							    pipeline_test_casse_name == "QAPipelineTests"
							    and tokenizer_name is not None
							    and not tokenizer_name.endswith("""Fast"""							)
							):
									# `QAPipelineTests` fails for a few models when the slower tokenizer are used.
									# (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
									# TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
									return True

							return False




					def 						snake_case_		(    self							)    ->	int:
							UpperCamelCase__								=						TFGPTJModelTester(self							)
							UpperCamelCase__								=						ConfigTester(self     ,   config_class=__lowerCAmelCase     ,   n_embd=37							)




					def 						snake_case_		(    self							)    ->	str:
							self.config_tester.run_common_tests()




					def 						snake_case_		(    self							)    ->	List[Any]:
							UpperCamelCase__								=						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_gptj_model(*__lowerCAmelCase							)




					def 						snake_case_		(    self							)    ->	Optional[Any]:
							UpperCamelCase__								=						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_gptj_model_past(*__lowerCAmelCase							)




					def 						snake_case_		(    self							)    ->	Optional[Any]:
							UpperCamelCase__								=						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_gptj_model_attention_mask_past(*__lowerCAmelCase							)




					def 						snake_case_		(    self							)    ->	List[str]:
							UpperCamelCase__								=						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_gptj_model_past_large_inputs(*__lowerCAmelCase							)




					def 						snake_case_		(    self							)    ->	Optional[Any]:
							UpperCamelCase__								=						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_gptj_lm_head_model(*__lowerCAmelCase							)




					@slow
					@unittest.skipIf(
					    not is_tf_available() or len(tf.config.list_physical_devices("""GPU"""							)							) > 0     ,   """skip testing on GPU for now to avoid GPU OOM."""     ,   )
					def 						snake_case_		(    self							)    ->	Tuple:
							UpperCamelCase__								=						TFGPTJModel.from_pretrained("""EleutherAI/gpt-j-6B"""     ,   from_pt=__lowerCAmelCase							)
							self.assertIsNotNone(__lowerCAmelCase							)




					@unittest.skip(reason="""Currently, model embeddings are going to undergo a major refactor."""							)
					def 						snake_case_		(    self							)    ->	Optional[Any]:
							super().test_resize_token_embeddings()




@require_tf
@tooslow
# Marked as @tooslow due to GPU OOM -- but still useful to run locally. Requires ~39GB of RAM.
class 	_UpperCAmelCase						( unittest.TestCase       ):




					def 						snake_case_		(    self							)    ->	Any:
							UpperCamelCase__								=						TFGPTJForCausalLM.from_pretrained("""EleutherAI/gpt-j-6B"""     ,   from_pt=__lowerCAmelCase							)
							UpperCamelCase__								=						tf.convert_to_tensor([[464, 3_290]]     ,   dtype=tf.intaa							)  # The dog
							# fmt: off
							# The dog is a man's best friend. It is a loyal companion, and it is a friend
							UpperCamelCase__								=						[464, 3_290, 318, 257, 582, 338, 1_266, 1_545, 13, 632, 318, 257, 9_112, 15_185, 11, 290, 340, 318, 257, 1_545]
							# fmt: on
							UpperCamelCase__								=						model.generate(__lowerCAmelCase     ,   do_sample=__lowerCAmelCase							)
							self.assertListEqual(output_ids[0].numpy().tolist()     ,   __lowerCAmelCase							)




					def 						snake_case_		(    self							)    ->	Any:
							UpperCamelCase__								=						AutoTokenizer.from_pretrained("""EleutherAI/gpt-j-6B"""     ,   revision="""float16"""							)
							UpperCamelCase__								=						TFGPTJForCausalLM.from_pretrained("""EleutherAI/gpt-j-6B"""     ,   revision="""float16"""     ,   from_pt=__lowerCAmelCase							)

							UpperCamelCase__								=						tokenizer("""Today is a nice day and"""     ,   return_tensors="""tf"""							)
							# forces the generation to happen on CPU, to avoid GPU-related quirks
							with tf.device(""":/CPU:0"""							):
									UpperCamelCase__								=						model.generate(**__lowerCAmelCase     ,   do_sample=__lowerCAmelCase     ,   seed=[42, 0]							)
							UpperCamelCase__								=						tokenizer.decode(output_ids[0]     ,   skip_special_tokens=__lowerCAmelCase							)

							UpperCamelCase__								=						"""Today is a nice day and I’m going to go for a walk. I’"""
							self.assertEqual(__lowerCAmelCase     ,   __lowerCAmelCase							)




					def 						snake_case_		(    self							)    ->	Optional[Any]:
							UpperCamelCase__								=						TFGPTJForCausalLM.from_pretrained("""EleutherAI/gpt-j-6B"""     ,   revision="""float16"""     ,   from_pt=__lowerCAmelCase							)
							UpperCamelCase__								=						AutoTokenizer.from_pretrained("""EleutherAI/gpt-j-6B"""     ,   revision="""float16"""							)

							UpperCamelCase__								=						"""left"""

							# Define PAD Token = EOS Token = 50256
							UpperCamelCase__								=						tokenizer.eos_token
							UpperCamelCase__								=						model.config.eos_token_id

							# use different length sentences to test batching
							UpperCamelCase__								=						[
							    """Hello, my dog is a little""",
							    """Today, I""",
							]
							UpperCamelCase__								=						[
							    """Hello, my dog is a little over a year old and has been diagnosed with hip dysplasia""",
							    """Today, I’m going to be talking about a topic that’""",
							]
							return model, tokenizer, sentences, expected_output_sentences




					def 						snake_case_		(    self							)    ->	int:
							# Confirms that we get the expected results with left-padded beam search
							UpperCamelCase__	,      UpperCamelCase__	,      UpperCamelCase__	,      UpperCamelCase__								=						self._get_beam_search_test_objects()

							UpperCamelCase__								=						tokenizer(__lowerCAmelCase     ,   return_tensors="""tf"""     ,   padding=__lowerCAmelCase							)
							UpperCamelCase__								=						model.generate(**__lowerCAmelCase     ,   do_sample=__lowerCAmelCase     ,   num_beams=2							)
							UpperCamelCase__								=						tokenizer.batch_decode(__lowerCAmelCase     ,   skip_special_tokens=__lowerCAmelCase							)
							self.assertListEqual(__lowerCAmelCase     ,   __lowerCAmelCase							)




					def 						snake_case_		(    self							)    ->	Optional[Any]:
							# Confirms that left-padding is working properly
							UpperCamelCase__	,      UpperCamelCase__	,      UpperCamelCase__	,      UpperCamelCase__								=						self._get_beam_search_test_objects()

							UpperCamelCase__								=						tokenizer(__lowerCAmelCase     ,   return_tensors="""tf"""     ,   padding=__lowerCAmelCase							)
							UpperCamelCase__								=						tokenizer(sentences[0]     ,   return_tensors="""tf"""							)
							UpperCamelCase__								=						model.generate(**__lowerCAmelCase     ,   do_sample=__lowerCAmelCase     ,   num_beams=2							)
							UpperCamelCase__								=						(
							    shape_list(inputs_non_padded["""input_ids"""]							)[-1]
							    - tf.reduce_sum(tf.cast(inputs["""attention_mask"""][-1]     ,   tf.intaa							)							).numpy()
							)
							UpperCamelCase__								=						tokenizer(sentences[1]     ,   return_tensors="""tf"""							)
							UpperCamelCase__								=						model.generate(
							    **__lowerCAmelCase     ,   do_sample=__lowerCAmelCase     ,   num_beams=2     ,   max_length=model.config.max_length - num_paddings							)
							UpperCamelCase__								=						tokenizer.decode(output_non_padded[0]     ,   skip_special_tokens=__lowerCAmelCase							)
							UpperCamelCase__								=						tokenizer.decode(output_padded[0]     ,   skip_special_tokens=__lowerCAmelCase							)
							self.assertListEqual(__lowerCAmelCase     ,   [non_padded_sentence, padded_sentence]							)




					def 						snake_case_		(    self							)    ->	List[str]:
							# Confirms that XLA is working properly
							UpperCamelCase__	,      UpperCamelCase__	,      UpperCamelCase__	,      UpperCamelCase__								=						self._get_beam_search_test_objects()

							UpperCamelCase__								=						tokenizer(__lowerCAmelCase     ,   return_tensors="""tf"""     ,   padding=__lowerCAmelCase							)
							UpperCamelCase__								=						tf.function(model.generate     ,   jit_compile=__lowerCAmelCase							)
							UpperCamelCase__								=						xla_generate(**__lowerCAmelCase     ,   do_sample=__lowerCAmelCase     ,   num_beams=2							)
							UpperCamelCase__								=						tokenizer.batch_decode(__lowerCAmelCase     ,   skip_special_tokens=__lowerCAmelCase							)
							self.assertListEqual(__lowerCAmelCase     ,   __lowerCAmelCase							)












import argparse

import torch

from diffusers import UNetaDConditionModel

def      a    (   _lowercase :  Dict			,    _lowercase :  Dict			,    _lowercase :  Any			,    _lowercase :  Tuple=None			,    _lowercase :  List[Any]=None			,    _lowercase :  Optional[Any]=None							)  ->Any:
		assert isinstance(_lowercase			,    _lowercase							), "Paths should be a list of dicts containing 'old' and 'new' keys."

		# Splits the attention layers into three variables.
		if attention_paths_to_split is not None:
				for path, path_map in attention_paths_to_split.items():
						UpperCamelCase__								=						old_checkpoint[path]
						UpperCamelCase__								=						old_tensor.shape[0] // 3

						UpperCamelCase__								=						(-1, channels) if len(old_tensor.shape							) == 3 else (-1)

						UpperCamelCase__								=						old_tensor.shape[0] // config["""num_head_channels"""] // 3

						UpperCamelCase__								=						old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]							)
						UpperCamelCase__	,      UpperCamelCase__	,      UpperCamelCase__								=						old_tensor.split(channels // num_heads			,    dim=1							)

						UpperCamelCase__								=						query.reshape(_lowercase							)
						UpperCamelCase__								=						key.reshape(_lowercase							)
						UpperCamelCase__								=						value.reshape(_lowercase							)

		for path in paths:
				UpperCamelCase__								=						path["""new"""]

				# These have already been assigned
				if attention_paths_to_split is not None and new_path in attention_paths_to_split:
						continue

				if additional_replacements is not None:
						for replacement in additional_replacements:
								UpperCamelCase__								=						new_path.replace(replacement["""old"""]			,    replacement["""new"""]							)

        # proj_attn.weight has to be converted from conv 1D to linear
				UpperCamelCase__								=						old_checkpoint[path["""old"""]]
				UpperCamelCase__								=						["""proj_attn.weight"""]
				UpperCamelCase__								=						["""proj_out.weight""", """proj_in.weight"""]
				if any(k in new_path for k in names							):
						UpperCamelCase__								=						weight[:, :, 0]
				elif any(k in new_path for k in names_a							) and len(weight.shape							) > 2 and ".attentions." not in new_path:
						UpperCamelCase__								=						weight[:, :, 0]
				else:
						UpperCamelCase__								=						weight

def      a    (   _lowercase :  Tuple			,    _lowercase :  List[Any]=0							)  ->int:
		UpperCamelCase__								=						[]
		for old_item in old_list:
				UpperCamelCase__								=						old_item

				#         new_item = new_item.replace('norm.weight', 'group_norm.weight')
				#         new_item = new_item.replace('norm.bias', 'group_norm.bias')

				#         new_item = new_item.replace('proj_out.weight', 'proj_attn.weight')
				#         new_item = new_item.replace('proj_out.bias', 'proj_attn.bias')

				#         new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)

				mapping.append({"""old""": old_item, """new""": new_item}							)

		return mapping

def      a    (   _lowercase :  Union[str, Any]			,    _lowercase :  int=1							)  ->str:
		if n_shave_prefix_segments >= 0:
				return ".".join(path.split("""."""							)[n_shave_prefix_segments:]							)
		else:
				return ".".join(path.split("""."""							)[:n_shave_prefix_segments]							)

def      a    (   _lowercase :  List[str]			,    _lowercase :  Any=0							)  ->Any:
		UpperCamelCase__								=						[]
		for old_item in old_list:
				mapping.append({"""old""": old_item, """new""": old_item}							)

		return mapping

def      a    (   _lowercase :  List[str]			,    _lowercase :  int=0							)  ->Tuple:
		UpperCamelCase__								=						[]
		for old_item in old_list:
				UpperCamelCase__								=						old_item.replace("""in_layers.0"""			,    """norm1"""							)
				UpperCamelCase__								=						new_item.replace("""in_layers.2"""			,    """conv1"""							)

				UpperCamelCase__								=						new_item.replace("""out_layers.0"""			,    """norm2"""							)
				UpperCamelCase__								=						new_item.replace("""out_layers.3"""			,    """conv2"""							)

				UpperCamelCase__								=						new_item.replace("""emb_layers.1"""			,    """time_emb_proj"""							)
				UpperCamelCase__								=						new_item.replace("""skip_connection"""			,    """conv_shortcut"""							)

				UpperCamelCase__								=						shave_segments(_lowercase			,    n_shave_prefix_segments=_lowercase							)

				if "temopral_conv" not in old_item:
						mapping.append({"""old""": old_item, """new""": new_item}							)

		return mapping

def      a    (   _lowercase :  List[Any]			,    _lowercase :  List[Any]			,    _lowercase :  int=None			,    _lowercase :  Optional[Any]=False							)  ->List[str]:
		UpperCamelCase__								=						{}
		UpperCamelCase__								=						list(checkpoint.keys()							)

		UpperCamelCase__								=						"""model.diffusion_model."""

		# at least a 100 parameters have to start with `model_ema` in order for the checkpoint to be EMA
		if sum(k.startswith("""model_ema"""							) for k in keys							) > 1_0_0 and extract_ema:
				print(F"""Checkpoint {path} has both EMA and non-EMA weights."""							)
				print(
				    """In this conversion only the EMA weights are extracted. If you want to instead extract the non-EMA"""
				    """ weights (useful to continue fine-tuning), please make sure to remove the `--extract_ema` flag."""							)
				for key in keys:
						if key.startswith("""model.diffusion_model"""							):
								UpperCamelCase__								=						"""model_ema.""" + """""".join(key.split("""."""							)[1:]							)
								UpperCamelCase__								=						checkpoint.pop(_lowercase							)
		else:
				if sum(k.startswith("""model_ema"""							) for k in keys							) > 1_0_0:
						print(
						    """In this conversion only the non-EMA weights are extracted. If you want to instead extract the EMA"""
						    """ weights (usually better for inference), please make sure to add the `--extract_ema` flag."""							)

				for key in keys:
						UpperCamelCase__								=						checkpoint.pop(_lowercase							)

		UpperCamelCase__								=						{}

		UpperCamelCase__								=						unet_state_dict["""time_embed.0.weight"""]
		UpperCamelCase__								=						unet_state_dict["""time_embed.0.bias"""]
		UpperCamelCase__								=						unet_state_dict["""time_embed.2.weight"""]
		UpperCamelCase__								=						unet_state_dict["""time_embed.2.bias"""]

		if config["class_embed_type"] is None:
				# No parameters to port
				...
		elif config["class_embed_type"] == "timestep" or config["class_embed_type"] == "projection":
				UpperCamelCase__								=						unet_state_dict["""label_emb.0.0.weight"""]
				UpperCamelCase__								=						unet_state_dict["""label_emb.0.0.bias"""]
				UpperCamelCase__								=						unet_state_dict["""label_emb.0.2.weight"""]
				UpperCamelCase__								=						unet_state_dict["""label_emb.0.2.bias"""]
		else:
				raise NotImplementedError(F"""Not implemented `class_embed_type`: {config['class_embed_type']}"""							)

		UpperCamelCase__								=						unet_state_dict["""input_blocks.0.0.weight"""]
		UpperCamelCase__								=						unet_state_dict["""input_blocks.0.0.bias"""]

		UpperCamelCase__								=						[v for v in unet_state_dict if v.startswith("""input_blocks.0.1"""							)]
		UpperCamelCase__								=						renew_attention_paths(_lowercase							)
		UpperCamelCase__								=						{"""old""": """input_blocks.0.1""", """new""": """transformer_in"""}
		assign_to_checkpoint(_lowercase			,    _lowercase			,    _lowercase			,    additional_replacements=[meta_path]			,    config=_lowercase							)

		UpperCamelCase__								=						unet_state_dict["""out.0.weight"""]
		UpperCamelCase__								=						unet_state_dict["""out.0.bias"""]
		UpperCamelCase__								=						unet_state_dict["""out.2.weight"""]
		UpperCamelCase__								=						unet_state_dict["""out.2.bias"""]

		# Retrieves the keys for the input blocks only
		UpperCamelCase__								=						len({""".""".join(layer.split("""."""							)[:2]							) for layer in unet_state_dict if """input_blocks""" in layer}							)
		UpperCamelCase__								=						{
		    layer_id: [key for key in unet_state_dict if F"""input_blocks.{layer_id}""" in key]
		    for layer_id in range(_lowercase							)
		}

		# Retrieves the keys for the middle blocks only
		UpperCamelCase__								=						len({""".""".join(layer.split("""."""							)[:2]							) for layer in unet_state_dict if """middle_block""" in layer}							)
		UpperCamelCase__								=						{
		    layer_id: [key for key in unet_state_dict if F"""middle_block.{layer_id}""" in key]
		    for layer_id in range(_lowercase							)
		}

		# Retrieves the keys for the output blocks only
		UpperCamelCase__								=						len({""".""".join(layer.split("""."""							)[:2]							) for layer in unet_state_dict if """output_blocks""" in layer}							)
		UpperCamelCase__								=						{
		    layer_id: [key for key in unet_state_dict if F"""output_blocks.{layer_id}""" in key]
		    for layer_id in range(_lowercase							)
		}

		for i in range(1			,    _lowercase							):
				UpperCamelCase__								=						(i - 1) // (config["""layers_per_block"""] + 1)
				UpperCamelCase__								=						(i - 1) % (config["""layers_per_block"""] + 1)

				UpperCamelCase__								=						[
				    key for key in input_blocks[i] if F"""input_blocks.{i}.0""" in key and F"""input_blocks.{i}.0.op""" not in key
				]
				UpperCamelCase__								=						[key for key in input_blocks[i] if F"""input_blocks.{i}.1""" in key]
				UpperCamelCase__								=						[key for key in input_blocks[i] if F"""input_blocks.{i}.2""" in key]

				if F"""input_blocks.{i}.op.weight""" in unet_state_dict:
						UpperCamelCase__								=						unet_state_dict.pop(
						    F"""input_blocks.{i}.op.weight"""							)
						UpperCamelCase__								=						unet_state_dict.pop(
						    F"""input_blocks.{i}.op.bias"""							)

				UpperCamelCase__								=						renew_resnet_paths(_lowercase							)
				UpperCamelCase__								=						{"""old""": F"""input_blocks.{i}.0""", """new""": F"""down_blocks.{block_id}.resnets.{layer_in_block_id}"""}
				assign_to_checkpoint(
				    _lowercase			,    _lowercase			,    _lowercase			,    additional_replacements=[meta_path]			,    config=_lowercase							)

				UpperCamelCase__								=						[key for key in resnets if """temopral_conv""" in key]
				UpperCamelCase__								=						renew_temp_conv_paths(_lowercase							)
				UpperCamelCase__								=						{
				    """old""": F"""input_blocks.{i}.0.temopral_conv""",
				    """new""": F"""down_blocks.{block_id}.temp_convs.{layer_in_block_id}""",
				}
				assign_to_checkpoint(
				    _lowercase			,    _lowercase			,    _lowercase			,    additional_replacements=[meta_path]			,    config=_lowercase							)

				if len(_lowercase							):
						UpperCamelCase__								=						renew_attention_paths(_lowercase							)
						UpperCamelCase__								=						{"""old""": F"""input_blocks.{i}.1""", """new""": F"""down_blocks.{block_id}.attentions.{layer_in_block_id}"""}
						assign_to_checkpoint(
						    _lowercase			,    _lowercase			,    _lowercase			,    additional_replacements=[meta_path]			,    config=_lowercase							)

				if len(_lowercase							):
						UpperCamelCase__								=						renew_attention_paths(_lowercase							)
						UpperCamelCase__								=						{
						    """old""": F"""input_blocks.{i}.2""",
						    """new""": F"""down_blocks.{block_id}.temp_attentions.{layer_in_block_id}""",
						}
						assign_to_checkpoint(
						    _lowercase			,    _lowercase			,    _lowercase			,    additional_replacements=[meta_path]			,    config=_lowercase							)

		UpperCamelCase__								=						middle_blocks[0]
		UpperCamelCase__								=						[key for key in resnet_a if """temopral_conv""" in key]
		UpperCamelCase__								=						middle_blocks[1]
		UpperCamelCase__								=						middle_blocks[2]
		UpperCamelCase__								=						middle_blocks[3]
		UpperCamelCase__								=						[key for key in resnet_a if """temopral_conv""" in key]

		UpperCamelCase__								=						renew_resnet_paths(_lowercase							)
		UpperCamelCase__								=						{"""old""": """middle_block.0""", """new""": """mid_block.resnets.0"""}
		assign_to_checkpoint(
		    _lowercase			,    _lowercase			,    _lowercase			,    config=_lowercase			,    additional_replacements=[meta_path]							)

		UpperCamelCase__								=						renew_temp_conv_paths(_lowercase							)
		UpperCamelCase__								=						{"""old""": """middle_block.0.temopral_conv""", """new""": """mid_block.temp_convs.0"""}
		assign_to_checkpoint(
		    _lowercase			,    _lowercase			,    _lowercase			,    config=_lowercase			,    additional_replacements=[meta_path]							)

		UpperCamelCase__								=						renew_resnet_paths(_lowercase							)
		UpperCamelCase__								=						{"""old""": """middle_block.3""", """new""": """mid_block.resnets.1"""}
		assign_to_checkpoint(
		    _lowercase			,    _lowercase			,    _lowercase			,    config=_lowercase			,    additional_replacements=[meta_path]							)

		UpperCamelCase__								=						renew_temp_conv_paths(_lowercase							)
		UpperCamelCase__								=						{"""old""": """middle_block.3.temopral_conv""", """new""": """mid_block.temp_convs.1"""}
		assign_to_checkpoint(
		    _lowercase			,    _lowercase			,    _lowercase			,    config=_lowercase			,    additional_replacements=[meta_path]							)

		UpperCamelCase__								=						renew_attention_paths(_lowercase							)
		UpperCamelCase__								=						{"""old""": """middle_block.1""", """new""": """mid_block.attentions.0"""}
		assign_to_checkpoint(
		    _lowercase			,    _lowercase			,    _lowercase			,    additional_replacements=[meta_path]			,    config=_lowercase							)

		UpperCamelCase__								=						renew_attention_paths(_lowercase							)
		UpperCamelCase__								=						{"""old""": """middle_block.2""", """new""": """mid_block.temp_attentions.0"""}
		assign_to_checkpoint(
		    _lowercase			,    _lowercase			,    _lowercase			,    additional_replacements=[meta_path]			,    config=_lowercase							)

		for i in range(_lowercase							):
				UpperCamelCase__								=						i // (config["""layers_per_block"""] + 1)
				UpperCamelCase__								=						i % (config["""layers_per_block"""] + 1)
				UpperCamelCase__								=						[shave_segments(_lowercase			,    2							) for name in output_blocks[i]]
				UpperCamelCase__								=						{}

				for layer in output_block_layers:
						UpperCamelCase__	,      UpperCamelCase__								=						layer.split("""."""							)[0], shave_segments(_lowercase			,    1							)
						if layer_id in output_block_list:
								output_block_list[layer_id].append(_lowercase							)
						else:
								UpperCamelCase__								=						[layer_name]

				if len(_lowercase							) > 1:
						UpperCamelCase__								=						[key for key in output_blocks[i] if F"""output_blocks.{i}.0""" in key]
						UpperCamelCase__								=						[key for key in output_blocks[i] if F"""output_blocks.{i}.1""" in key]
						UpperCamelCase__								=						[key for key in output_blocks[i] if F"""output_blocks.{i}.2""" in key]

						UpperCamelCase__								=						renew_resnet_paths(_lowercase							)
						UpperCamelCase__								=						renew_resnet_paths(_lowercase							)

						UpperCamelCase__								=						{"""old""": F"""output_blocks.{i}.0""", """new""": F"""up_blocks.{block_id}.resnets.{layer_in_block_id}"""}
						assign_to_checkpoint(
						    _lowercase			,    _lowercase			,    _lowercase			,    additional_replacements=[meta_path]			,    config=_lowercase							)

						UpperCamelCase__								=						[key for key in resnets if """temopral_conv""" in key]
						UpperCamelCase__								=						renew_temp_conv_paths(_lowercase							)
						UpperCamelCase__								=						{
						    """old""": F"""output_blocks.{i}.0.temopral_conv""",
						    """new""": F"""up_blocks.{block_id}.temp_convs.{layer_in_block_id}""",
						}
						assign_to_checkpoint(
						    _lowercase			,    _lowercase			,    _lowercase			,    additional_replacements=[meta_path]			,    config=_lowercase							)

						UpperCamelCase__								=						{k: sorted(_lowercase							) for k, v in output_block_list.items()}
						if ["conv.bias", "conv.weight"] in output_block_list.values():
								UpperCamelCase__								=						list(output_block_list.values()							).index(["""conv.bias""", """conv.weight"""]							)
								UpperCamelCase__								=						unet_state_dict[
								    F"""output_blocks.{i}.{index}.conv.weight"""
								]
								UpperCamelCase__								=						unet_state_dict[
								    F"""output_blocks.{i}.{index}.conv.bias"""
								]

								# Clear attentions as they have been attributed above.
								if len(_lowercase							) == 2:
										UpperCamelCase__								=						[]

						if len(_lowercase							):
								UpperCamelCase__								=						renew_attention_paths(_lowercase							)
								UpperCamelCase__								=						{
								    """old""": F"""output_blocks.{i}.1""",
								    """new""": F"""up_blocks.{block_id}.attentions.{layer_in_block_id}""",
								}
								assign_to_checkpoint(
								    _lowercase			,    _lowercase			,    _lowercase			,    additional_replacements=[meta_path]			,    config=_lowercase							)

						if len(_lowercase							):
								UpperCamelCase__								=						renew_attention_paths(_lowercase							)
								UpperCamelCase__								=						{
								    """old""": F"""output_blocks.{i}.2""",
								    """new""": F"""up_blocks.{block_id}.temp_attentions.{layer_in_block_id}""",
								}
								assign_to_checkpoint(
								    _lowercase			,    _lowercase			,    _lowercase			,    additional_replacements=[meta_path]			,    config=_lowercase							)
				else:
						UpperCamelCase__								=						renew_resnet_paths(_lowercase			,    n_shave_prefix_segments=1							)
						for path in resnet_0_paths:
								UpperCamelCase__								=						""".""".join(["""output_blocks""", str(_lowercase							), path["""old"""]]							)
								UpperCamelCase__								=						""".""".join(["""up_blocks""", str(_lowercase							), """resnets""", str(_lowercase							), path["""new"""]]							)
								UpperCamelCase__								=						unet_state_dict[old_path]

						UpperCamelCase__								=						[l for l in output_block_layers if """temopral_conv""" in l]
						for path in temopral_conv_paths:
								UpperCamelCase__								=						path.split("""temopral_conv."""							)[-1]
								UpperCamelCase__								=						""".""".join(["""output_blocks""", str(_lowercase							), str(_lowercase							), """temopral_conv""", pruned_path]							)
								UpperCamelCase__								=						""".""".join(["""up_blocks""", str(_lowercase							), """temp_convs""", str(_lowercase							), pruned_path]							)
								UpperCamelCase__								=						unet_state_dict[old_path]

		return new_checkpoint


if __name__ == "__main__":
				__lowerCAmelCase 	=      argparse.ArgumentParser()

				parser.add_argument(
				    "--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert."
				)
				parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.")
				__lowerCAmelCase 	=      parser.parse_args()

				__lowerCAmelCase 	=      torch.load(args.checkpoint_path, map_location="cpu")
				__lowerCAmelCase 	=      UNetaDConditionModel()

				__lowerCAmelCase 	=      convert_ldm_unet_checkpoint(unet_checkpoint, unet.config)

				__lowerCAmelCase 	=      set(unet.state_dict().keys()) - set(converted_ckpt.keys())
				__lowerCAmelCase 	=      set(converted_ckpt.keys()) - set(unet.state_dict().keys())

				assert len(diff_a) == len(diff_a) == 0, "Converted weights don't match"

				# load state_dict
				unet.load_state_dict(converted_ckpt)

				unet.save_pretrained(args.dump_path)

				# -- finish converting the unet --





