



import gc
import math
import unittest

import torch

from diffusers import UNetaDModel
from diffusers.utils import floats_tensor, logging, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


__UpperCAmelCase			  =			logging.get_logger(__name__)

enable_full_determinism()



class 							lowerCAmelCase_						(    a__			,       a__			,       unittest.TestCase				):
       UpperCAmelCase__							:   int  	=    UNetaDModel
       UpperCAmelCase__							:   List[Any]  	=    "sample"
       @property
       def    snake_case_			(     self			)	->						Optional[Any]:
          UpperCamelCase       :    List[Any]		      =      4
          UpperCamelCase       :    Any		      =      3
          UpperCamelCase       :    List[Any]		      =      (32, 32)

          UpperCamelCase       :    Union[str, Any]		      =      floats_tensor((batch_size, num_channels) + sizes			).to(SCREAMING_SNAKE_CASE_			)
          UpperCamelCase       :    Any		      =      torch.tensor([10]			).to(SCREAMING_SNAKE_CASE_			)

          return {"sample": noise, "timestep": time_step}
       @property
       def    snake_case_			(     self			)	->						List[str]:
          return (3, 32, 32)
       @property
       def    snake_case_			(     self			)	->						str:
          return (3, 32, 32)






       def    snake_case_			(     self			)	->						Any:
          UpperCamelCase       :    Dict		      =      {
              'block_out_channels': (32, 64),
              'down_block_types': ('DownBlock2D', 'AttnDownBlock2D'),
              'up_block_types': ('AttnUpBlock2D', 'UpBlock2D'),
              'attention_head_dim': 3,
              'out_channels': 3,
              'in_channels': 3,
              'layers_per_block': 2,
              'sample_size': 32,
          }
          UpperCamelCase       :    Tuple		      =      self.dummy_input
          return init_dict, inputs_dict






class 							lowerCAmelCase_						(    a__			,       a__			,       unittest.TestCase				):
       UpperCAmelCase__							:   Optional[int]  	=    UNetaDModel
       UpperCAmelCase__							:   Optional[int]  	=    "sample"
       @property
       def    snake_case_			(     self			)	->						Any:
          UpperCamelCase       :    List[str]		      =      4
          UpperCamelCase       :    Tuple		      =      4
          UpperCamelCase       :    Tuple		      =      (32, 32)

          UpperCamelCase       :    Optional[int]		      =      floats_tensor((batch_size, num_channels) + sizes			).to(SCREAMING_SNAKE_CASE_			)
          UpperCamelCase       :    List[Any]		      =      torch.tensor([10]			).to(SCREAMING_SNAKE_CASE_			)

          return {"sample": noise, "timestep": time_step}
       @property
       def    snake_case_			(     self			)	->						Optional[int]:
          return (4, 32, 32)
       @property
       def    snake_case_			(     self			)	->						Union[str, Any]:
          return (4, 32, 32)
       def    snake_case_			(     self			)	->						str:
          UpperCamelCase       :    Union[str, Any]		      =      {
              'sample_size': 32,
              'in_channels': 4,
              'out_channels': 4,
              'layers_per_block': 2,
              'block_out_channels': (32, 64),
              'attention_head_dim': 32,
              'down_block_types': ('DownBlock2D', 'DownBlock2D'),
              'up_block_types': ('UpBlock2D', 'UpBlock2D'),
          }
          UpperCamelCase       :    Optional[Any]		      =      self.dummy_input
          return init_dict, inputs_dict
       def    snake_case_			(     self			)	->						Optional[Any]:
          UpperCamelCase      ,	UpperCamelCase       :    Any		      =      UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update',     output_loading_info=SCREAMING_SNAKE_CASE_			)

          self.assertIsNotNone(SCREAMING_SNAKE_CASE_			)
          self.assertEqual(len(loading_info['missing_keys']			),     0			)

          model.to(SCREAMING_SNAKE_CASE_			)
          UpperCamelCase       :    Optional[int]		      =      model(**self.dummy_input			).sample

          assert image is not None, "Make sure output is not None"
       @unittest.skipIf(torch_device != 'cuda',     'This test is supposed to run on GPU'			)
       def    snake_case_			(     self			)	->						Any:
          UpperCamelCase      ,	UpperCamelCase       :    int		      =      UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update',     output_loading_info=SCREAMING_SNAKE_CASE_			)
          model.to(SCREAMING_SNAKE_CASE_			)
          UpperCamelCase       :    List[str]		      =      model(**self.dummy_input			).sample

          assert image is not None, "Make sure output is not None"
       @unittest.skipIf(torch_device != 'cuda',     'This test is supposed to run on GPU'			)
       def    snake_case_			(     self			)	->						str:
          # by defautl model loading will use accelerate as `low_cpu_mem_usage=True`
          UpperCamelCase      ,	UpperCamelCase       :    Optional[int]		      =      UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update',     output_loading_info=SCREAMING_SNAKE_CASE_			)
          model_accelerate.to(SCREAMING_SNAKE_CASE_			)
          model_accelerate.eval()

          UpperCamelCase       :    str		      =      torch.randn(
              1,     model_accelerate.config.in_channels,     model_accelerate.config.sample_size,     model_accelerate.config.sample_size,     generator=torch.manual_seed(0			),     )
          UpperCamelCase       :    Union[str, Any]		      =      noise.to(SCREAMING_SNAKE_CASE_			)
          UpperCamelCase       :    Optional[Any]		      =      torch.tensor([10] * noise.shape[0]			).to(SCREAMING_SNAKE_CASE_			)

          UpperCamelCase       :    List[str]		      =      model_accelerate(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			)['sample']

          # two models don't need to stay in the device at the same time
          del model_accelerate
          torch.cuda.empty_cache()
          gc.collect()

          UpperCamelCase      ,	UpperCamelCase       :    Optional[Any]		      =      UNetaDModel.from_pretrained(
              'fusing/unet-ldm-dummy-update',     output_loading_info=SCREAMING_SNAKE_CASE_,     low_cpu_mem_usage=SCREAMING_SNAKE_CASE_			)
          model_normal_load.to(SCREAMING_SNAKE_CASE_			)
          model_normal_load.eval()
          UpperCamelCase       :    Tuple		      =      model_normal_load(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			)['sample']

          assert torch_all_close(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     rtol=1e-3			)






       def    snake_case_			(     self			)	->						List[str]:
          UpperCamelCase       :    Union[str, Any]		      =      UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'			)
          model.eval()
          model.to(SCREAMING_SNAKE_CASE_			)

          UpperCamelCase       :    Dict		      =      torch.randn(
              1,     model.config.in_channels,     model.config.sample_size,     model.config.sample_size,     generator=torch.manual_seed(0			),     )
          UpperCamelCase       :    Optional[Any]		      =      noise.to(SCREAMING_SNAKE_CASE_			)
          UpperCamelCase       :    List[Any]		      =      torch.tensor([10] * noise.shape[0]			).to(SCREAMING_SNAKE_CASE_			)

          with torch.no_grad():
             UpperCamelCase       :    List[str]		      =      model(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			).sample

          UpperCamelCase       :    int		      =      output[0, -1, -3:, -3:].flatten().cpu()
          # fmt: off
          UpperCamelCase       :    Tuple		      =      torch.tensor([-13.32_58, -20.11_00, -15.98_73, -17.66_17, -23.05_96, -17.94_19, -13.36_75, -16.18_89, -12.38_00]			)
          # fmt: on

          self.assertTrue(torch_all_close(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     rtol=1e-3			)			)






class 							lowerCAmelCase_						(    a__			,       a__			,       unittest.TestCase				):
       UpperCAmelCase__							:   str  	=    UNetaDModel
       UpperCAmelCase__							:   Optional[Any]  	=    "sample"
       @property
       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_=(32, 32)			)	->						Any:
          UpperCamelCase       :    List[Any]		      =      4
          UpperCamelCase       :    str		      =      3

          UpperCamelCase       :    Tuple		      =      floats_tensor((batch_size, num_channels) + sizes			).to(SCREAMING_SNAKE_CASE_			)
          UpperCamelCase       :    int		      =      torch.tensor(batch_size * [10]			).to(dtype=torch.intaa,     device=SCREAMING_SNAKE_CASE_			)

          return {"sample": noise, "timestep": time_step}
       @property
       def    snake_case_			(     self			)	->						str:
          return (3, 32, 32)
       @property
       def    snake_case_			(     self			)	->						Optional[Any]:
          return (3, 32, 32)
       def    snake_case_			(     self			)	->						Tuple:
          UpperCamelCase       :    List[str]		      =      {
              'block_out_channels': [32, 64, 64, 64],
              'in_channels': 3,
              'layers_per_block': 1,
              'out_channels': 3,
              'time_embedding_type': 'fourier',
              'norm_eps': 1e-6,
              'mid_block_scale_factor': math.sqrt(2.0			),
              'norm_num_groups': None,
              'down_block_types': [
                  'SkipDownBlock2D',
                  'AttnSkipDownBlock2D',
                  'SkipDownBlock2D',
                  'SkipDownBlock2D',
              ],
              'up_block_types': [
                  'SkipUpBlock2D',
                  'SkipUpBlock2D',
                  'AttnSkipUpBlock2D',
                  'SkipUpBlock2D',
              ],
          }
          UpperCamelCase       :    Union[str, Any]		      =      self.dummy_input
          return init_dict, inputs_dict
       @slow
       def    snake_case_			(     self			)	->						Any:
          UpperCamelCase      ,	UpperCamelCase       :    Tuple		      =      UNetaDModel.from_pretrained('google/ncsnpp-celebahq-256',     output_loading_info=SCREAMING_SNAKE_CASE_			)
          self.assertIsNotNone(SCREAMING_SNAKE_CASE_			)
          self.assertEqual(len(loading_info['missing_keys']			),     0			)

          model.to(SCREAMING_SNAKE_CASE_			)
          UpperCamelCase       :    int		      =      self.dummy_input
          UpperCamelCase       :    Optional[Any]		      =      floats_tensor((4, 3) + (256, 256)			).to(SCREAMING_SNAKE_CASE_			)
          UpperCamelCase       :    Optional[Any]		      =      noise
          UpperCamelCase       :    Union[str, Any]		      =      model(**SCREAMING_SNAKE_CASE_			)

          assert image is not None, "Make sure output is not None"
       @slow
       def    snake_case_			(     self			)	->						List[str]:
          UpperCamelCase       :    List[Any]		      =      UNetaDModel.from_pretrained('google/ncsnpp-celebahq-256'			)
          model.to(SCREAMING_SNAKE_CASE_			)

          UpperCamelCase       :    Union[str, Any]		      =      4
          UpperCamelCase       :    Any		      =      3
          UpperCamelCase       :    Union[str, Any]		      =      (256, 256)

          UpperCamelCase       :    Optional[int]		      =      torch.ones((batch_size, num_channels) + sizes			).to(SCREAMING_SNAKE_CASE_			)
          UpperCamelCase       :    Union[str, Any]		      =      torch.tensor(batch_size * [1e-4]			).to(SCREAMING_SNAKE_CASE_			)

          with torch.no_grad():
             UpperCamelCase       :    Optional[int]		      =      model(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			).sample

          UpperCamelCase       :    str		      =      output[0, -3:, -3:, -1].flatten().cpu()
          # fmt: off
          UpperCamelCase       :    Tuple		      =      torch.tensor([-48_42.86_91, -64_99.66_31, -38_00.19_53, -79_78.26_86, -1_09_80.71_29, -2_00_28.85_35, 81_48.28_22, 23_42.29_05, 5_67.76_08]			)
          # fmt: on

          self.assertTrue(torch_all_close(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     rtol=1e-2			)			)
       def    snake_case_			(     self			)	->						Optional[int]:
          UpperCamelCase       :    Tuple		      =      UNetaDModel.from_pretrained('fusing/ncsnpp-ffhq-ve-dummy-update'			)
          model.to(SCREAMING_SNAKE_CASE_			)

          UpperCamelCase       :    str		      =      4
          UpperCamelCase       :    Dict		      =      3
          UpperCamelCase       :    str		      =      (32, 32)

          UpperCamelCase       :    int		      =      torch.ones((batch_size, num_channels) + sizes			).to(SCREAMING_SNAKE_CASE_			)
          UpperCamelCase       :    int		      =      torch.tensor(batch_size * [1e-4]			).to(SCREAMING_SNAKE_CASE_			)

          with torch.no_grad():
             UpperCamelCase       :    Any		      =      model(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			).sample

          UpperCamelCase       :    Union[str, Any]		      =      output[0, -3:, -3:, -1].flatten().cpu()
          # fmt: off
          UpperCamelCase       :    Tuple		      =      torch.tensor([-0.03_25, -0.09_00, -0.08_69, -0.03_32, -0.07_25, -0.02_70, -0.01_01, 0.02_27, 0.02_56]			)
          # fmt: on

          self.assertTrue(torch_all_close(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     rtol=1e-2			)			)






       def    snake_case_			(     self			)	->						Optional[int]:
          # not required for this model
          pass




"""simple docstring"""



import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


SCREAMING_SNAKE_CASE__:List[str]									=			logging.getLogger(__name__)



def 	_lowerCamelCase(	):
       __a       =    argparse.ArgumentParser(
           description="Prepare TFRecord shards from pre-tokenized samples of the wikitext dataset."				)
       parser.add_argument(
           "--dataset_name"   ,    type=a   ,    default="wikitext"   ,    help="Name of the training. Explore datasets at: hf.co/datasets."   ,    )
       parser.add_argument(
           "--dataset_config"   ,    type=a   ,    default="wikitext-103-raw-v1"   ,    help="Configuration name of the dataset."				)
       parser.add_argument(
           "--tokenizer_name_or_path"   ,    type=a   ,    default="sayakpaul/unigram-tokenizer-wikitext"   ,    help="Tokenizer identifier. Can be a local filepath or a Hub identifier."   ,    )
       parser.add_argument(
           "--shard_size"   ,    type=a   ,    default=1_0_0_0   ,    help="Number of entries to go in a single shard."   ,    )
       parser.add_argument("--split"   ,    type=a   ,    default="train"   ,    choices=["train", "test", "validation"]				)
       parser.add_argument(
           "--limit"   ,    default=a   ,    type=a   ,    help="Limit the number of shards (used for debugging)."   ,    )
       parser.add_argument(
           "--max_length"   ,    type=a   ,    default=5_1_2   ,    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=a   ,    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."   ,    )

       __a       =    parser.parse_args()
       return args



def 	_lowerCamelCase(	a				):
       def fn(a				):
              return tokenizer(examples["text"]				)

       return fn



def 	_lowerCamelCase(	a				):
       __a       =    []
       for i in range(len(tokenized_data["input_ids"]				)				):
              __a       =    {
                  "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]				)				),
              }
              __a       =    tf.train.Features(feature=a				)
              __a       =    tf.train.Example(features=a				)
              __a       =    example.SerializeToString()
              records.append(a				)
       return records



def 	_lowerCamelCase(	a				):
       __a       =    datasets.load_dataset(args.dataset_name   ,    args.dataset_config   ,    split=args.split				)

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

       __a       =    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				)
              __a       =    os.path.join(args.output_dir   ,    args.split				)
              if not os.path.exists(a				):
                     os.makedirs(a				)
       else:
              __a       =    os.path.join(args.output_dir   ,    args.split				)

       # Tokenize the whole dataset at once.
       __a       =    tokenize_function(a				)
       __a       =    dataset.map(a   ,    batched=a   ,    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(a				):
              # Concatenate all texts.
              __a       =    {k: sum(examples[k]   ,    []				) for k in examples.keys()}
              __a       =    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 🫀
              __a       =    (total_length // args.max_length) * args.max_length
              # Split by chunks of max_len.
              __a       =    {
                  k: [t[i : i + args.max_length] for i in range(0   ,    a   ,    args.max_length				)]
                  for k, t in concatenated_examples.items()
              }
              return result

       __a       =    dataset_tokenized.map(a   ,    batched=a   ,    batch_size=1_0_0_0   ,    num_proc=4				)

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

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

              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=a				)


if __name__ == "__main__":
   SCREAMING_SNAKE_CASE__:List[Any]									=			parse_args()
   main(args)
