


'''simple docstring'''



import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


A__		:					Any												=  logging.getLogger(__name__)



def 			UpperCAmelCase__					(	)     ->		Dict:
     __lowerCamelCase   :	str								=		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.'      ,    )

     __lowerCamelCase   :	Union[str, Any]								=		parser.parse_args()
     return args



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				Optional[Any]							)     ->		Optional[int]:
     def fn(UpperCAmelCase_    :				int							):
          return tokenizer(examples['text']							)

     return fn



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				int							)     ->		Tuple:
     __lowerCamelCase   :	int								=		[]
     for i in range(len(tokenized_data['input_ids']							)							):
          __lowerCamelCase   :	Optional[int]								=		{
              '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]							)							),
          }
          __lowerCamelCase   :	Union[str, Any]								=		tf.train.Features(feature=UpperCAmelCase_							)
          __lowerCamelCase   :	Any								=		tf.train.Example(features=UpperCAmelCase_							)
          __lowerCamelCase   :	List[Any]								=		example.SerializeToString()
          records.append(UpperCAmelCase_							)
     return records



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				Tuple							)     ->		int:
     __lowerCamelCase   :	Any								=		datasets.load_dataset(args.dataset_name      ,    args.dataset_config      ,    split=args.split							)

     if args.limit is not None:
          __lowerCamelCase   :	Tuple								=		min(len(UpperCAmelCase_							)      ,    args.limit							)
          __lowerCamelCase   :	Union[str, Any]								=		dataset.select(range(UpperCAmelCase_							)							)
          print(F'Limiting the dataset to {args.limit} entries.'							)

     __lowerCamelCase   :	Any								=		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							)
          __lowerCamelCase   :	Any								=		os.path.join(args.output_dir      ,    args.split							)
          if not os.path.exists(UpperCAmelCase_							):
               os.makedirs(UpperCAmelCase_							)
     else:
          __lowerCamelCase   :	Optional[Any]								=		os.path.join(args.output_dir      ,    args.split							)

     # Tokenize the whole dataset at once.
     __lowerCamelCase   :	Union[str, Any]								=		tokenize_function(UpperCAmelCase_							)
     __lowerCamelCase   :	Optional[int]								=		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_    :				Optional[int]							):
          # Concatenate all texts.
          __lowerCamelCase   :	List[Any]								=		{k: sum(examples[k]      ,    []							) for k in examples.keys()}
          __lowerCamelCase   :	int								=		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 🫀
          __lowerCamelCase   :	List[Any]								=		(total_length // args.max_length) * args.max_length
          # Split by chunks of max_len.
          __lowerCamelCase   :	Any								=		{
              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

     __lowerCamelCase   :	Union[str, Any]								=		dataset_tokenized.map(UpperCAmelCase_      ,    batched=UpperCAmelCase_      ,    batch_size=10_00      ,    num_proc=4							)

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

          with tf.io.TFRecordWriter(UpperCAmelCase_							) as out_file:
               for i in range(len(UpperCAmelCase_							)							):
                    __lowerCamelCase   :	Union[str, Any]								=		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__":
      A__		:					int												=  parse_args()
      main(args)










'''simple docstring'''



import unittest

import torch

from diffusers import VQModel
from diffusers.utils import floats_tensor, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()





class    UpperCAmelCase_  (_UpperCAmelCase					,		_UpperCAmelCase					,		unittest.TestCase  ):





    """simple docstring"""






    lowerCamelCase  :					Union[str, Any]							     =			VQModel
    lowerCamelCase  :					Union[str, Any]							     =			'sample'







    @property
    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_=(32, 32)	)				->     Any:
         __lowerCamelCase   :	Tuple								=		4
         __lowerCamelCase   :	Union[str, Any]								=		3

         __lowerCamelCase   :	List[Any]								=		floats_tensor((batch_size, num_channels) + sizes	).to(SCREAMING_SNAKE_CASE_	)

         return {"sample": image}







    @property
    def       lowercase_				(							self	)				->     Optional[int]:
         return (3, 32, 32)







    @property
    def       lowercase_				(							self	)				->     List[Any]:
         return (3, 32, 32)







    def       lowercase_				(							self	)				->     Optional[int]:
         __lowerCamelCase   :	Optional[Any]								=		{
             'block_out_channels': [32, 64],
             'in_channels': 3,
             'out_channels': 3,
             'down_block_types': ['DownEncoderBlock2D', 'DownEncoderBlock2D'],
             'up_block_types': ['UpDecoderBlock2D', 'UpDecoderBlock2D'],
             'latent_channels': 3,
         }
         __lowerCamelCase   :	Union[str, Any]								=		self.dummy_input
         return init_dict, inputs_dict







    def       lowercase_				(							self	)				->     str:
         pass







    def       lowercase_				(							self	)				->     Optional[int]:
         pass







    def       lowercase_				(							self	)				->     Any:
         __lowerCamelCase				,       __lowerCamelCase   :	Union[str, Any]								=		VQModel.from_pretrained('fusing/vqgan-dummy'			,     output_loading_info=SCREAMING_SNAKE_CASE_	)
         self.assertIsNotNone(SCREAMING_SNAKE_CASE_	)
         self.assertEqual(len(loading_info['missing_keys']	)			,     0	)

         model.to(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	List[Any]								=		model(**self.dummy_input	)

         assert image is not None, "Make sure output is not None"







    def       lowercase_				(							self	)				->     int:
         __lowerCamelCase   :	List[Any]								=		VQModel.from_pretrained('fusing/vqgan-dummy'	)
         model.to(SCREAMING_SNAKE_CASE_	).eval()

         torch.manual_seed(0	)
         if torch.cuda.is_available():
              torch.cuda.manual_seed_all(0	)

         __lowerCamelCase   :	List[str]								=		torch.randn(1			,     model.config.in_channels			,     model.config.sample_size			,     model.config.sample_size	)
         __lowerCamelCase   :	Optional[int]								=		image.to(SCREAMING_SNAKE_CASE_	)
         with torch.no_grad():
              __lowerCamelCase   :	int								=		model(SCREAMING_SNAKE_CASE_	).sample

         __lowerCamelCase   :	List[str]								=		output[0, -1, -3:, -3:].flatten().cpu()
         # fmt: off
         __lowerCamelCase   :	Union[str, Any]								=		torch.tensor([-0.0_1_5_3, -0.4_0_4_4, -0.1_8_8_0, -0.5_1_6_1, -0.2_4_1_8, -0.4_0_7_2, -0.1_6_1_2, -0.0_6_3_3, -0.0_1_4_3]	)
         # fmt: on
         self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     atol=1E-3	)	)






