






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 				_UpperCamelCase      :


      def __init__(						self       ,    __UpperCamelCase       ,    __UpperCamelCase=3       ,    __UpperCamelCase=7       ,    __UpperCamelCase=True       ,    __UpperCamelCase=True       ,    __UpperCamelCase=False       ,    __UpperCamelCase=True       ,    __UpperCamelCase=9_9       ,    __UpperCamelCase=3_2       ,    __UpperCamelCase=5       ,    __UpperCamelCase=4       ,    __UpperCamelCase=3_7       ,    __UpperCamelCase="gelu"       ,    __UpperCamelCase=0.1       ,    __UpperCamelCase=0.1       ,    __UpperCamelCase=5_1_2       ,    __UpperCamelCase=1_6       ,    __UpperCamelCase=2       ,    __UpperCamelCase=0.0_2       ,    __UpperCamelCase=3       ,    __UpperCamelCase=4       ,    __UpperCamelCase=None       ,    )->						int:
           __lowerCAmelCase             =       parent
           __lowerCAmelCase             =       batch_size
           __lowerCAmelCase             =       seq_length
           __lowerCAmelCase             =       is_training
           __lowerCAmelCase             =       use_input_mask
           __lowerCAmelCase             =       use_token_type_ids
           __lowerCAmelCase             =       use_labels
           __lowerCAmelCase             =       vocab_size
           __lowerCAmelCase             =       hidden_size
           __lowerCAmelCase             =       num_hidden_layers
           __lowerCAmelCase             =       num_attention_heads
           __lowerCAmelCase             =       intermediate_size
           __lowerCAmelCase             =       hidden_act
           __lowerCAmelCase             =       hidden_dropout_prob
           __lowerCAmelCase             =       attention_probs_dropout_prob
           __lowerCAmelCase             =       max_position_embeddings
           __lowerCAmelCase             =       type_vocab_size
           __lowerCAmelCase             =       type_sequence_label_size
           __lowerCAmelCase             =       initializer_range
           __lowerCAmelCase             =       num_labels
           __lowerCAmelCase             =       num_choices
           __lowerCAmelCase             =       scope


      def 			__UpperCAmelCase		(						self						)->						Optional[Any]:
           __lowerCAmelCase             =       ids_tensor([self.batch_size, self.seq_length]       ,    self.vocab_size						)

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

           __lowerCAmelCase             =       None

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

           __lowerCAmelCase             =       self.get_config()

           return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels


      def 			__UpperCAmelCase		(						self						)->						Dict:
           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=__a       ,    initializer_range=self.initializer_range       ,    pad_token_id=1       ,    new_decoder_architecture=__a       ,    )


      def 			__UpperCAmelCase		(						self       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase						)->						Any:
           __lowerCAmelCase             =       FalconModel(config=__a						)
           model.to(__a						)
           model.eval()
           __lowerCAmelCase             =       model(__a       ,    attention_mask=__a						)
           __lowerCAmelCase             =       model(__a						)
           self.parent.assertEqual(result.last_hidden_state.shape       ,    (self.batch_size, self.seq_length, self.hidden_size)						)


      def 			__UpperCAmelCase		(						self       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    )->						int:
           __lowerCAmelCase             =       True
           __lowerCAmelCase             =       FalconModel(__a						)
           model.to(__a						)
           model.eval()
           __lowerCAmelCase             =       model(
               __a       ,    attention_mask=__a       ,    encoder_hidden_states=__a       ,    encoder_attention_mask=__a       ,    )
           __lowerCAmelCase             =       model(
               __a       ,    attention_mask=__a       ,    encoder_hidden_states=__a       ,    )
           __lowerCAmelCase             =       model(__a       ,    attention_mask=__a						)
           self.parent.assertEqual(result.last_hidden_state.shape       ,    (self.batch_size, self.seq_length, self.hidden_size)						)


      def 			__UpperCAmelCase		(						self       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    )->						Optional[Any]:
           __lowerCAmelCase             =       FalconForCausalLM(config=__a						)
           model.to(__a						)
           model.eval()
           __lowerCAmelCase             =       model(__a       ,    attention_mask=__a       ,    labels=__a						)
           self.parent.assertEqual(result.logits.shape       ,    (self.batch_size, self.seq_length, self.vocab_size)						)


      def 			__UpperCAmelCase		(						self       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    )->						List[Any]:
           __lowerCAmelCase             =       True
           __lowerCAmelCase             =       True
           __lowerCAmelCase             =       FalconForCausalLM(config=__a						)
           model.to(__a						)
           model.eval()

           # first forward pass
           __lowerCAmelCase             =       model(
               __a       ,    attention_mask=__a       ,    encoder_hidden_states=__a       ,    encoder_attention_mask=__a       ,    use_cache=__a       ,    )
           __lowerCAmelCase             =       outputs.past_key_values

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

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

           __lowerCAmelCase             =       model(
               __a       ,    attention_mask=__a       ,    encoder_hidden_states=__a       ,    encoder_attention_mask=__a       ,    output_hidden_states=__a       ,    )["""hidden_states"""][0]
           __lowerCAmelCase             =       model(
               __a       ,    attention_mask=__a       ,    encoder_hidden_states=__a       ,    encoder_attention_mask=__a       ,    past_key_values=__a       ,    output_hidden_states=__a       ,    )["""hidden_states"""][0]

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


      def 			__UpperCAmelCase		(						self						)->						Optional[Any]:
           __lowerCAmelCase             =       self.prepare_config_and_inputs()
           (
               __lowerCAmelCase	
           )             =       config_and_inputs
           __lowerCAmelCase             =       {"""input_ids""": input_ids, """attention_mask""": input_mask}
           return config, inputs_dict



@require_torch
class 				_UpperCamelCase      (a__				,		a__				,		a__				,		unittest.TestCase      ):
      snake_case_	=				(
          (
              FalconModel,
              FalconForCausalLM,
              FalconForSequenceClassification,
              FalconForTokenClassification,
              FalconForQuestionAnswering,
          )
          if is_torch_available()
          else ()
      )
      snake_case_	=				(FalconForCausalLM,) if is_torch_available() else ()
      snake_case_	=				(
          {
              """feature-extraction""": FalconModel,
              """text-classification""": FalconForSequenceClassification,
              """text-generation""": FalconForCausalLM,
              """question-answering""": FalconForQuestionAnswering,
              """token-classification""": FalconForTokenClassification,
              """zero-shot""": FalconForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      snake_case_	=				False
      snake_case_	=				False


      def 			__UpperCAmelCase		(						self						)->						Optional[Any]:
           __lowerCAmelCase             =       FalconModelTester(self						)
           __lowerCAmelCase             =       ConfigTester(self       ,    config_class=__a       ,    hidden_size=3_7						)


      def 			__UpperCAmelCase		(						self						)->						Dict:
           self.config_tester.run_common_tests()


      def 			__UpperCAmelCase		(						self						)->						int:
           __lowerCAmelCase             =       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*__a						)


      def 			__UpperCAmelCase		(						self						)->						Optional[int]:
           __lowerCAmelCase             =       self.model_tester.prepare_config_and_inputs()
           for alibi in [True, False]:
                __lowerCAmelCase             =       alibi
                self.model_tester.create_and_check_model(__a       ,    *__a						)


      def 			__UpperCAmelCase		(						self						)->						List[Any]:
           __lowerCAmelCase             =       self.model_tester.prepare_config_and_inputs_for_common()
           __lowerCAmelCase             =       3
           __lowerCAmelCase             =       input_dict["""input_ids"""]
           __lowerCAmelCase             =       input_ids.ne(1						).to(__a						)
           __lowerCAmelCase             =       ids_tensor([self.model_tester.batch_size]       ,    self.model_tester.type_sequence_label_size						)
           __lowerCAmelCase             =       FalconForSequenceClassification(__a						)
           model.to(__a						)
           model.eval()
           __lowerCAmelCase             =       model(__a       ,    attention_mask=__a       ,    labels=__a						)
           self.assertEqual(result.logits.shape       ,    (self.model_tester.batch_size, self.model_tester.num_labels)						)


      def 			__UpperCAmelCase		(						self						)->						Tuple:
           __lowerCAmelCase             =       self.model_tester.prepare_config_and_inputs_for_common()
           __lowerCAmelCase             =       3
           __lowerCAmelCase             =       """single_label_classification"""
           __lowerCAmelCase             =       input_dict["""input_ids"""]
           __lowerCAmelCase             =       input_ids.ne(1						).to(__a						)
           __lowerCAmelCase             =       ids_tensor([self.model_tester.batch_size]       ,    self.model_tester.type_sequence_label_size						)
           __lowerCAmelCase             =       FalconForSequenceClassification(__a						)
           model.to(__a						)
           model.eval()
           __lowerCAmelCase             =       model(__a       ,    attention_mask=__a       ,    labels=__a						)
           self.assertEqual(result.logits.shape       ,    (self.model_tester.batch_size, self.model_tester.num_labels)						)


      def 			__UpperCAmelCase		(						self						)->						List[str]:
           __lowerCAmelCase             =       self.model_tester.prepare_config_and_inputs_for_common()
           __lowerCAmelCase             =       input_dict["""input_ids"""]
           __lowerCAmelCase             =       FalconForCausalLM(__a						)
           model.to(__a						)
           model.eval()
           __lowerCAmelCase             =       model(__a       ,    use_cache=__a						)
           __lowerCAmelCase             =       input_ids.shape[0]
           __lowerCAmelCase             =       model._convert_to_rw_cache(result.past_key_values						)
           __lowerCAmelCase             =       model._convert_cache_to_standard_format(__a       ,    __a						)
           for layer in range(len(__a						)						):
                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 			__UpperCAmelCase		(						self						)->						Dict:
           __lowerCAmelCase             =       self.model_tester.prepare_config_and_inputs_for_common()
           __lowerCAmelCase             =       3
           __lowerCAmelCase             =       """multi_label_classification"""
           __lowerCAmelCase             =       input_dict["""input_ids"""]
           __lowerCAmelCase             =       input_ids.ne(1						).to(__a						)
           __lowerCAmelCase             =       ids_tensor(
               [self.model_tester.batch_size, config.num_labels]       ,    self.model_tester.type_sequence_label_size						).to(torch.float						)
           __lowerCAmelCase             =       FalconForSequenceClassification(__a						)
           model.to(__a						)
           model.eval()
           __lowerCAmelCase             =       model(__a       ,    attention_mask=__a       ,    labels=__a						)
           self.assertEqual(result.logits.shape       ,    (self.model_tester.batch_size, self.model_tester.num_labels)						)


      def 			__UpperCAmelCase		(						self						)->						str:
           # 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             =       self.model_tester.prepare_config_and_inputs_for_common()

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

                __lowerCAmelCase             =       model_class(__a						).to(__a						)
                if "use_cache" not in inputs:
                     __lowerCAmelCase             =       True
                __lowerCAmelCase             =       model(**__a						)

                # 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             =       (
                    getattr(__a       ,    "decoder_layers"       ,    __a						)
                    or getattr(__a       ,    "num_decoder_layers"       ,    __a						)
                    or config.num_hidden_layers
                )
                __lowerCAmelCase             =       getattr(__a       ,    "num_kv_heads"       ,    config.num_attention_heads						)
                __lowerCAmelCase             =       getattr(__a       ,    "d_model"       ,    config.hidden_size						)
                __lowerCAmelCase             =       embed_dim // num_attention_heads

                __lowerCAmelCase             =       outputs["""past_key_values"""]
                self.assertEqual(len(__a						)       ,    __a						)

                __lowerCAmelCase             =       inputs["""input_ids"""].shape
                for i in range(__a						):
                     if config.new_decoder_architecture:
                          __lowerCAmelCase             =       config.num_attention_heads
                     elif config.multi_query:
                          __lowerCAmelCase             =       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 				_UpperCamelCase      (unittest.TestCase      ):


      @slow
      def 			__UpperCAmelCase		(						self						)->						Any:
           __lowerCAmelCase             =       AutoTokenizer.from_pretrained("Rocketknight1/falcon-rw-1b"						)
           __lowerCAmelCase             =       FalconForCausalLM.from_pretrained("Rocketknight1/falcon-rw-1b"						)
           model.eval()
           model.to(__a						)
           __lowerCAmelCase             =       tokenizer("My favorite food is"       ,    return_tensors="pt"						).to(__a						)

           __lowerCAmelCase             =       (
               """My favorite food is pizza. I love it so much that I have a pizza party every year for my birthday."""
           )

           __lowerCAmelCase             =       model.generate(**__a       ,    do_sample=__a       ,    max_new_tokens=1_9						)
           __lowerCAmelCase             =       tokenizer.batch_decode(__a						)[0]

           self.assertEqual(__a       ,    __a						)


      @slow
      def 			__UpperCAmelCase		(						self						)->						List[Any]:
           # 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             =       AutoTokenizer.from_pretrained(__a						)
                __lowerCAmelCase             =       FalconForCausalLM.from_pretrained(__a						)
                model.eval()
                model.to(__a						)
                __lowerCAmelCase             =       tokenizer("My favorite food is"       ,    return_tensors="pt"						).to(__a						)

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


      @slow
      def 			__UpperCAmelCase		(						self						)->						Tuple:
           # 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             =       AutoTokenizer.from_pretrained(__a						)
                     __lowerCAmelCase             =       FalconForCausalLM.from_pretrained(__a						)
                     model.eval()
                     model.to(device=__a						)
                     __lowerCAmelCase             =       tokenizer("My favorite food is"       ,    return_tensors="pt"						).to(__a						)

                     # Test results are the same with and without cache
                     __lowerCAmelCase             =       model.generate(**__a       ,    do_sample=__a       ,    max_new_tokens=2_0       ,    use_cache=__a						)
                     __lowerCAmelCase             =       model.generate(**__a       ,    do_sample=__a       ,    max_new_tokens=2_0       ,    use_cache=__a						)
                     self.assertTrue((outputs_cache - outputs_no_cache).sum().item() == 0						)










import argparse

from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection

from diffusers import UnCLIPImageVariationPipeline, UnCLIPPipeline


if __name__ == "__main__":
      snake_case__				:	Optional[int]					       =					argparse.ArgumentParser()

      parser.add_argument('--dump_path', default=None, type=str, required=True, help='Path to the output model.')

      parser.add_argument(
          '--txt2img_unclip',
          default='kakaobrain/karlo-v1-alpha',
          type=str,
          required=False,
          help='The pretrained txt2img unclip.',
      )

      snake_case__				:	Tuple					       =					parser.parse_args()

      snake_case__				:	Dict					       =					UnCLIPPipeline.from_pretrained(args.txtaimg_unclip)

      snake_case__				:	Optional[int]					       =					CLIPImageProcessor()
      snake_case__				:	Optional[Any]					       =					CLIPVisionModelWithProjection.from_pretrained('openai/clip-vit-large-patch14')

      snake_case__				:	Dict					       =					UnCLIPImageVariationPipeline(
          decoder=txtaimg.decoder,
          text_encoder=txtaimg.text_encoder,
          tokenizer=txtaimg.tokenizer,
          text_proj=txtaimg.text_proj,
          feature_extractor=feature_extractor,
          image_encoder=image_encoder,
          super_res_first=txtaimg.super_res_first,
          super_res_last=txtaimg.super_res_last,
          decoder_scheduler=txtaimg.decoder_scheduler,
          super_res_scheduler=txtaimg.super_res_scheduler,
      )

      imgaimg.save_pretrained(args.dump_path)





