

import unittest

import numpy as np
import torch
from torch import nn
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import KandinskyVaaPriorPipeline, PriorTransformer, UnCLIPScheduler
from diffusers.utils import torch_device
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()
class      SCREAMING_SNAKE_CASE__  (    _lowerCAmelCase			,						unittest.TestCase		):





       '''simple docstring'''




       UpperCamelCase__			:  Optional[Any]					     =							KandinskyVaaPriorPipeline
       UpperCamelCase__			:  List[Any]					     =							['''prompt''']
       UpperCamelCase__			:  List[Any]					     =							['''prompt''', '''negative_prompt''']
       UpperCamelCase__			:  Dict					     =							[
           '''num_images_per_prompt''',
           '''generator''',
           '''num_inference_steps''',
           '''latents''',
           '''negative_prompt''',
           '''guidance_scale''',
           '''output_type''',
           '''return_dict''',
       ]
       UpperCamelCase__			:  List[Any]					     =							False

       @property
       def UpperCAmelCase_    (    self							:    Union[str, Any]				)      ->     List[str]:
             return 32

       @property
       def UpperCAmelCase_    (    self							:    Dict				)      ->     int:
             return 32

       @property
       def UpperCAmelCase_    (    self							:    List[str]				)      ->     int:
             return self.time_input_dim

       @property
       def UpperCAmelCase_    (    self							:    Dict				)      ->     Dict:
             return self.time_input_dim * 4

       @property
       def UpperCAmelCase_    (    self							:    Optional[Any]				)      ->     Any:
             return 100

       @property
       def UpperCAmelCase_    (    self							:    Tuple				)      ->     Union[str, Any]:
             snake_case__											= CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""				)
             return tokenizer

       @property
       def UpperCAmelCase_    (    self							:    Union[str, Any]				)      ->     Optional[Any]:
             torch.manual_seed(0				)
             snake_case__											= CLIPTextConfig(
                 bos_token_id=0     ,		eos_token_id=2     ,		hidden_size=self.text_embedder_hidden_size     ,		projection_dim=self.text_embedder_hidden_size     ,		intermediate_size=37     ,		layer_norm_eps=1E-05     ,		num_attention_heads=4     ,		num_hidden_layers=5     ,		pad_token_id=1     ,		vocab_size=1000     ,		)
             return CLIPTextModelWithProjection(UpperCAmelCase_				)

       @property
       def UpperCAmelCase_    (    self							:    Optional[Any]				)      ->     str:
             torch.manual_seed(0				)

             snake_case__											= {
                 """num_attention_heads""": 2,
                 """attention_head_dim""": 12,
                 """embedding_dim""": self.text_embedder_hidden_size,
                 """num_layers""": 1,
             }

             snake_case__											= PriorTransformer(**UpperCAmelCase_				)
             # clip_std and clip_mean is initialized to be 0 so PriorTransformer.post_process_latents will always return 0 - set clip_std to be 1 so it won't return 0
             snake_case__											= nn.Parameter(torch.ones(model.clip_std.shape				)				)
             return model

       @property
       def UpperCAmelCase_    (    self							:    List[str]				)      ->     List[Any]:
             torch.manual_seed(0				)
             snake_case__											= CLIPVisionConfig(
                 hidden_size=self.text_embedder_hidden_size     ,		image_size=224     ,		projection_dim=self.text_embedder_hidden_size     ,		intermediate_size=37     ,		num_attention_heads=4     ,		num_channels=3     ,		num_hidden_layers=5     ,		patch_size=14     ,		)

             snake_case__											= CLIPVisionModelWithProjection(UpperCAmelCase_				)
             return model

       @property
       def UpperCAmelCase_    (    self							:    List[str]				)      ->     Union[str, Any]:
             snake_case__											= CLIPImageProcessor(
                 crop_size=224     ,		do_center_crop=UpperCAmelCase_     ,		do_normalize=UpperCAmelCase_     ,		do_resize=UpperCAmelCase_     ,		image_mean=[0.48_145_466, 0.4_578_275, 0.40_821_073]     ,		image_std=[0.26_862_954, 0.26_130_258, 0.27_577_711]     ,		resample=3     ,		size=224     ,		)

             return image_processor

       def UpperCAmelCase_    (    self							:    Optional[Any]				)      ->     Dict:
             snake_case__											= self.dummy_prior
             snake_case__											= self.dummy_image_encoder
             snake_case__											= self.dummy_text_encoder
             snake_case__											= self.dummy_tokenizer
             snake_case__											= self.dummy_image_processor

             snake_case__											= UnCLIPScheduler(
                 variance_type="""fixed_small_log"""     ,		prediction_type="""sample"""     ,		num_train_timesteps=1000     ,		clip_sample=UpperCAmelCase_     ,		clip_sample_range=10.0     ,		)

             snake_case__											= {
                 """prior""": prior,
                 """image_encoder""": image_encoder,
                 """text_encoder""": text_encoder,
                 """tokenizer""": tokenizer,
                 """scheduler""": scheduler,
                 """image_processor""": image_processor,
             }

             return components

       def UpperCAmelCase_    (    self							:    Any     ,		lowerCAmelCase__							:    Optional[Any]     ,		lowerCAmelCase__							:    List[str]=0				)      ->     Optional[Any]:
             if str(UpperCAmelCase_				).startswith("""mps"""				):
                   snake_case__											= torch.manual_seed(UpperCAmelCase_				)
             else:
                   snake_case__											= torch.Generator(device=UpperCAmelCase_				).manual_seed(UpperCAmelCase_				)
             snake_case__											= {
                 """prompt""": """horse""",
                 """generator""": generator,
                 """guidance_scale""": 4.0,
                 """num_inference_steps""": 2,
                 """output_type""": """np""",
             }
             return inputs

       def UpperCAmelCase_    (    self							:    int				)      ->     List[Any]:
             snake_case__											= """cpu"""

             snake_case__											= self.get_dummy_components()

             snake_case__											= self.pipeline_class(**UpperCAmelCase_				)
             snake_case__											= pipe.to(UpperCAmelCase_				)

             pipe.set_progress_bar_config(disable=UpperCAmelCase_				)

             snake_case__											= pipe(**self.get_dummy_inputs(UpperCAmelCase_				)				)
             snake_case__											= output.image_embeds

             snake_case__											= pipe(
                 **self.get_dummy_inputs(UpperCAmelCase_				)     ,		return_dict=UpperCAmelCase_     ,		)[0]

             snake_case__											= image[0, -10:]
             snake_case__											= image_from_tuple[0, -10:]

             assert image.shape == (1, 32)

             snake_case__											= np.array(
                 [-0.0_532, 1.7_120, 0.3_656, -1.0_852, -0.8_946, -1.1_756, 0.4_348, 0.2_482, 0.5_146, -0.1_156]				)

             assert np.abs(image_slice.flatten() - expected_slice				).max() < 1E-2
             assert np.abs(image_from_tuple_slice.flatten() - expected_slice				).max() < 1E-2

       @skip_mps
       def UpperCAmelCase_    (    self							:    List[Any]				)      ->     List[Any]:
             snake_case__											= torch_device == """cpu"""
             snake_case__											= True
             snake_case__											= False

             self._test_inference_batch_single_identical(
                 test_max_difference=UpperCAmelCase_     ,		relax_max_difference=UpperCAmelCase_     ,		test_mean_pixel_difference=UpperCAmelCase_     ,		)

       @skip_mps
       def UpperCAmelCase_    (    self							:    Dict				)      ->     Union[str, Any]:
             snake_case__											= torch_device == """cpu"""
             snake_case__											= False

             self._test_attention_slicing_forward_pass(
                 test_max_difference=UpperCAmelCase_     ,		test_mean_pixel_difference=UpperCAmelCase_     ,		)








"""simple docstring"""


import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


SCREAMING_SNAKE_CASE			  =  logging.getLogger(__name__)



def 						lowerCamelCase__      (				)->				Tuple:







       """simple docstring"""
       UpperCamelCase       			=		argparse.ArgumentParser(
           description="Prepare TFRecord shards from pre-tokenized samples of the wikitext dataset."					)
       parser.add_argument(
           "--dataset_name"		,   type=UpperCAmelCase_		,   default="wikitext"		,   help="Name of the training. Explore datasets at: hf.co/datasets."		,   )
       parser.add_argument(
           "--dataset_config"		,   type=UpperCAmelCase_		,   default="wikitext-103-raw-v1"		,   help="Configuration name of the dataset."					)
       parser.add_argument(
           "--tokenizer_name_or_path"		,   type=UpperCAmelCase_		,   default="sayakpaul/unigram-tokenizer-wikitext"		,   help="Tokenizer identifier. Can be a local filepath or a Hub identifier."		,   )
       parser.add_argument(
           "--shard_size"		,   type=UpperCAmelCase_		,   default=10_00		,   help="Number of entries to go in a single shard."		,   )
       parser.add_argument("--split"		,   type=UpperCAmelCase_		,   default="train"		,   choices=["train", "test", "validation"]					)
       parser.add_argument(
           "--limit"		,   default=UpperCAmelCase_		,   type=UpperCAmelCase_		,   help="Limit the number of shards (used for debugging)."		,   )
       parser.add_argument(
           "--max_length"		,   type=UpperCAmelCase_		,   default=5_12		,   help="Maximum sequence length. For training on TPUs, it helps to have a maximum"
           " sequence length that is a multiple of 8."		,   )
       parser.add_argument(
           "--output_dir"		,   default="tf-tpu"		,   type=UpperCAmelCase_		,   help="Output directory where the TFRecord shards will be saved. If the"
           " path is appended with `gs://` ('gs://tf-tpu', for example) then the TFRecord"
           " shards will be directly saved to a Google Cloud Storage bucket."		,   )

       UpperCamelCase       			=		parser.parse_args()
       return args



def 						lowerCamelCase__      (				UpperCAmelCase_					)->				Optional[int]:







       """simple docstring"""
       def fn(UpperCAmelCase_					):
              return tokenizer(examples["text"]					)

       return fn



def 						lowerCamelCase__      (				UpperCAmelCase_					)->				Union[str, Any]:







       """simple docstring"""
       UpperCamelCase       			=		[]
       for i in range(len(tokenized_data["input_ids"]					)					):
              UpperCamelCase       			=		{
                  "input_ids": tf.train.Feature(intaa_list=tf.train.IntaaList(value=tokenized_data["input_ids"][i]					)					),
                  "attention_mask": tf.train.Feature(
                      intaa_list=tf.train.IntaaList(value=tokenized_data["attention_mask"][i]					)					),
              }
              UpperCamelCase       			=		tf.train.Features(feature=UpperCAmelCase_					)
              UpperCamelCase       			=		tf.train.Example(features=UpperCAmelCase_					)
              UpperCamelCase       			=		example.SerializeToString()
              records.append(UpperCAmelCase_					)
       return records



def 						lowerCamelCase__      (				UpperCAmelCase_					)->				int:







       """simple docstring"""
       UpperCamelCase       			=		datasets.load_dataset(args.dataset_name		,   args.dataset_config		,   split=args.split					)

       if args.limit is not None:
              UpperCamelCase       			=		min(len(UpperCAmelCase_					)		,   args.limit					)
              UpperCamelCase       			=		dataset.select(range(UpperCAmelCase_					)					)
              print(F"Limiting the dataset to {args.limit} entries."					)

       UpperCamelCase       			=		AutoTokenizer.from_pretrained(args.tokenizer_name_or_path					)

       # Handle output directory creation.
       # For serializing into a Google Cloud Storage Bucket, one needs to first
       # create a bucket.
       if "gs" not in args.output_dir:
              if not os.path.exists(args.output_dir					):
                     os.makedirs(args.output_dir					)
              UpperCamelCase       			=		os.path.join(args.output_dir		,   args.split					)
              if not os.path.exists(UpperCAmelCase_					):
                     os.makedirs(UpperCAmelCase_					)
       else:
              UpperCamelCase       			=		os.path.join(args.output_dir		,   args.split					)

       # Tokenize the whole dataset at once.
       UpperCamelCase       			=		tokenize_function(UpperCAmelCase_					)
       UpperCamelCase       			=		dataset.map(UpperCAmelCase_		,   batched=UpperCAmelCase_		,   num_proc=4		,   remove_columns=["text"]					)

       # We need to concatenate all our texts together, and then split the result
       # into chunks of a fixed size, which we will call block_size. To do this, we
       # will use the map method again, with the option batched=True. When we use batched=True,
       # the function we pass to map() will be passed multiple inputs at once, allowing us
       # to group them into more or fewer examples than we had in the input.
       # This allows us to create our new fixed-length samples. The advantage of this
       # method is that we don't lose a whole lot of content from the dataset compared to the
       # case where we simply tokenize with a pre-defined max_length.

       def group_texts(UpperCAmelCase_					):
              # Concatenate all texts.
              UpperCamelCase       			=		{k: sum(examples[k]		,   []					) for k in examples.keys()}
              UpperCamelCase       			=		len(concatenated_examples[list(examples.keys()					)[0]]					)
              # We drop the small remainder, though you could add padding instead if the model supports it
              # In this, as in all things, we advise you to follow your heart 🫀
              UpperCamelCase       			=		(total_length // args.max_length) * args.max_length
              # Split by chunks of max_len.
              UpperCamelCase       			=		{
                  k: [t[i : i + args.max_length] for i in range(0		,   UpperCAmelCase_		,   args.max_length					)]
                  for k, t in concatenated_examples.items()
              }
              return result

       UpperCamelCase       			=		dataset_tokenized.map(UpperCAmelCase_		,   batched=UpperCAmelCase_		,   batch_size=10_00		,   num_proc=4					)

       UpperCamelCase       			=		0
       UpperCamelCase       			=		0
       for shard in range(0		,   len(UpperCAmelCase_					)		,   args.shard_size					):
              UpperCamelCase       			=		grouped_dataset[shard : shard + args.shard_size]
              UpperCamelCase       			=		len(dataset_snapshot["input_ids"]					)
              UpperCamelCase       			=		os.path.join(UpperCAmelCase_		,   F"dataset-{shard_count}-{records_containing}.tfrecord"					)
              UpperCamelCase       			=		get_serialized_examples(UpperCAmelCase_					)

              with tf.io.TFRecordWriter(UpperCAmelCase_					) as out_file:
                     for i in range(len(UpperCAmelCase_					)					):
                            UpperCamelCase       			=		serialized_examples[i]
                            out_file.write(UpperCAmelCase_					)
                     print("Wrote file {} containing {} records".format(UpperCAmelCase_		,   UpperCAmelCase_					)					)

              shard_count += 1
              total_records += records_containing

       with open(F"split-{args.split}-records-count.txt"		,   "w"					) as f:
              print(F"Total {args.split} records: {total_records}"		,   file=UpperCAmelCase_					)


if __name__ == "__main__":
      SCREAMING_SNAKE_CASE			  =  parse_args()
      main(args)





