


import gc
import unittest

import numpy as np
import torch
import torch.nn.functional as F
from transformers import (
    ClapTextConfig,
    ClapTextModelWithProjection,
    RobertaTokenizer,
    SpeechTaHifiGan,
    SpeechTaHifiGanConfig,
)

from diffusers import (
    AudioLDMPipeline,
    AutoencoderKL,
    DDIMScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.utils import is_xformers_available, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from ..pipeline_params import TEXT_TO_AUDIO_BATCH_PARAMS, TEXT_TO_AUDIO_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()





class   SCREAMING_SNAKE_CASE_			(a__						,       unittest.TestCase						):




       '''simple docstring'''



       _a   							=							AudioLDMPipeline
       _a   							=							TEXT_TO_AUDIO_PARAMS
       _a   							=							TEXT_TO_AUDIO_BATCH_PARAMS
       _a   							=							frozenset(
           [
               "num_inference_steps",
               "num_waveforms_per_prompt",
               "generator",
               "latents",
               "output_type",
               "return_dict",
               "callback",
               "callback_steps",
           ]						)



       def 				_lowerCAmelCase							(						self				:				Tuple							)					->Any:
          torch.manual_seed(0							)
          lowerCamelCase_						:   List[Any]		= UNetaDConditionModel(
              block_out_channels=(32, 64)		,				layers_per_block=2		,				sample_size=32		,				in_channels=4		,				out_channels=4		,				down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")		,				up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")		,				cross_attention_dim=(32, 64)		,				class_embed_type="""simple_projection"""		,				projection_class_embeddings_input_dim=32		,				class_embeddings_concat=__a		,				)
          lowerCamelCase_						:   Union[str, Any]		= DDIMScheduler(
              beta_start=0.00_085		,				beta_end=0.012		,				beta_schedule="""scaled_linear"""		,				clip_sample=__a		,				set_alpha_to_one=__a		,				)
          torch.manual_seed(0							)
          lowerCamelCase_						:   Optional[Any]		= AutoencoderKL(
              block_out_channels=[32, 64]		,				in_channels=1		,				out_channels=1		,				down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]		,				up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]		,				latent_channels=4		,				)
          torch.manual_seed(0							)
          lowerCamelCase_						:   Optional[Any]		= ClapTextConfig(
              bos_token_id=0		,				eos_token_id=2		,				hidden_size=32		,				intermediate_size=37		,				layer_norm_eps=1e-05		,				num_attention_heads=4		,				num_hidden_layers=5		,				pad_token_id=1		,				vocab_size=1_000		,				projection_dim=32		,				)
          lowerCamelCase_						:   Union[str, Any]		= ClapTextModelWithProjection(__a							)
          lowerCamelCase_						:   Optional[int]		= RobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-random-roberta"""		,				model_max_length=77							)

          lowerCamelCase_						:   Tuple		= SpeechTaHifiGanConfig(
              model_in_dim=8		,				sampling_rate=16_000		,				upsample_initial_channel=16		,				upsample_rates=[2, 2]		,				upsample_kernel_sizes=[4, 4]		,				resblock_kernel_sizes=[3, 7]		,				resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5]]		,				normalize_before=__a		,				)

          lowerCamelCase_						:   Tuple		= SpeechTaHifiGan(__a							)

          lowerCamelCase_						:   List[str]		= {
              """unet""": unet,
              """scheduler""": scheduler,
              """vae""": vae,
              """text_encoder""": text_encoder,
              """tokenizer""": tokenizer,
              """vocoder""": vocoder,
          }
          return components



       def 				_lowerCAmelCase							(						self				:				Tuple		,				__a				:				Tuple		,				__a				:				Optional[Any]=0							)					->List[Any]:
          if str(__a							).startswith("""mps"""							):
             lowerCamelCase_						:   List[str]		= torch.manual_seed(__a							)
          else:
             lowerCamelCase_						:   Any		= torch.Generator(device=__a							).manual_seed(__a							)
          lowerCamelCase_						:   Any		= {
              """prompt""": """A hammer hitting a wooden surface""",
              """generator""": generator,
              """num_inference_steps""": 2,
              """guidance_scale""": 6.0,
          }
          return inputs



       def 				_lowerCAmelCase							(						self				:				int							)					->Union[str, Any]:
          lowerCamelCase_						:   Dict		= """cpu"""  # ensure determinism for the device-dependent torch.Generator

          lowerCamelCase_						:   int		= self.get_dummy_components()
          lowerCamelCase_						:   Optional[int]		= AudioLDMPipeline(**__a							)
          lowerCamelCase_						:   str		= audioldm_pipe.to(__a							)
          audioldm_pipe.set_progress_bar_config(disable=__a							)

          lowerCamelCase_						:   Optional[int]		= self.get_dummy_inputs(__a							)
          lowerCamelCase_						:   int		= audioldm_pipe(**__a							)
          lowerCamelCase_						:   Dict		= output.audios[0]

          assert audio.ndim == 1
          assert len(__a							) == 256

          lowerCamelCase_						:   Any		= audio[:10]
          lowerCamelCase_						:   List[str]		= np.array(
              [-0.0_050, 0.0_050, -0.0_060, 0.0_033, -0.0_026, 0.0_033, -0.0_027, 0.0_033, -0.0_028, 0.0_033]							)

          assert np.abs(audio_slice - expected_slice							).max() < 1e-2



       def 				_lowerCAmelCase							(						self				:				Union[str, Any]							)					->Union[str, Any]:
          lowerCamelCase_						:   Tuple		= self.get_dummy_components()
          lowerCamelCase_						:   Any		= AudioLDMPipeline(**__a							)
          lowerCamelCase_						:   Optional[Any]		= audioldm_pipe.to(__a							)
          lowerCamelCase_						:   Union[str, Any]		= audioldm_pipe.to(__a							)
          audioldm_pipe.set_progress_bar_config(disable=__a							)

          lowerCamelCase_						:   str		= self.get_dummy_inputs(__a							)
          lowerCamelCase_						:   int		= 3 * [inputs["""prompt"""]]

          # forward
          lowerCamelCase_						:   List[Any]		= audioldm_pipe(**__a							)
          lowerCamelCase_						:   List[Any]		= output.audios[0]

          lowerCamelCase_						:   Any		= self.get_dummy_inputs(__a							)
          lowerCamelCase_						:   List[str]		= 3 * [inputs.pop("""prompt"""							)]

          lowerCamelCase_						:   int		= audioldm_pipe.tokenizer(
              __a		,				padding="""max_length"""		,				max_length=audioldm_pipe.tokenizer.model_max_length		,				truncation=__a		,				return_tensors="""pt"""		,				)
          lowerCamelCase_						:   str		= text_inputs["""input_ids"""].to(__a							)

          lowerCamelCase_						:   str		= audioldm_pipe.text_encoder(
              __a		,				)
          lowerCamelCase_						:   Dict		= prompt_embeds.text_embeds
          # additional L_2 normalization over each hidden-state
          lowerCamelCase_						:   str		= F.normalize(__a		,				dim=-1							)

          lowerCamelCase_						:   Tuple		= prompt_embeds

          # forward
          lowerCamelCase_						:   List[Any]		= audioldm_pipe(**__a							)
          lowerCamelCase_						:   Optional[Any]		= output.audios[0]

          assert np.abs(audio_a - audio_a							).max() < 1e-2



       def 				_lowerCAmelCase							(						self				:				Tuple							)					->Dict:
          lowerCamelCase_						:   List[Any]		= self.get_dummy_components()
          lowerCamelCase_						:   Optional[int]		= AudioLDMPipeline(**__a							)
          lowerCamelCase_						:   List[str]		= audioldm_pipe.to(__a							)
          lowerCamelCase_						:   str		= audioldm_pipe.to(__a							)
          audioldm_pipe.set_progress_bar_config(disable=__a							)

          lowerCamelCase_						:   Union[str, Any]		= self.get_dummy_inputs(__a							)
          lowerCamelCase_						:   str		= 3 * ["""this is a negative prompt"""]
          lowerCamelCase_						:   int		= negative_prompt
          lowerCamelCase_						:   Optional[Any]		= 3 * [inputs["""prompt"""]]

          # forward
          lowerCamelCase_						:   Optional[int]		= audioldm_pipe(**__a							)
          lowerCamelCase_						:   Optional[Any]		= output.audios[0]

          lowerCamelCase_						:   Optional[Any]		= self.get_dummy_inputs(__a							)
          lowerCamelCase_						:   str		= 3 * [inputs.pop("""prompt"""							)]

          lowerCamelCase_						:   str		= []
          for p in [prompt, negative_prompt]:
             lowerCamelCase_						:   List[str]		= audioldm_pipe.tokenizer(
                 __a		,				padding="""max_length"""		,				max_length=audioldm_pipe.tokenizer.model_max_length		,				truncation=__a		,				return_tensors="""pt"""		,				)
             lowerCamelCase_						:   Union[str, Any]		= text_inputs["""input_ids"""].to(__a							)

             lowerCamelCase_						:   Tuple		= audioldm_pipe.text_encoder(
                 __a		,				)
             lowerCamelCase_						:   List[Any]		= text_embeds.text_embeds
             # additional L_2 normalization over each hidden-state
             lowerCamelCase_						:   Dict		= F.normalize(__a		,				dim=-1							)

             embeds.append(__a							)

          lowerCamelCase_,	lowerCamelCase_						:   List[str]		= embeds

          # forward
          lowerCamelCase_						:   List[str]		= audioldm_pipe(**__a							)
          lowerCamelCase_						:   Any		= output.audios[0]

          assert np.abs(audio_a - audio_a							).max() < 1e-2



       def 				_lowerCAmelCase							(						self				:				Optional[Any]							)					->str:
          lowerCamelCase_						:   int		= """cpu"""  # ensure determinism for the device-dependent torch.Generator
          lowerCamelCase_						:   Tuple		= self.get_dummy_components()
          lowerCamelCase_						:   int		= PNDMScheduler(skip_prk_steps=__a							)
          lowerCamelCase_						:   Tuple		= AudioLDMPipeline(**__a							)
          lowerCamelCase_						:   Any		= audioldm_pipe.to(__a							)
          audioldm_pipe.set_progress_bar_config(disable=__a							)

          lowerCamelCase_						:   int		= self.get_dummy_inputs(__a							)
          lowerCamelCase_						:   int		= """egg cracking"""
          lowerCamelCase_						:   Tuple		= audioldm_pipe(**__a		,				negative_prompt=__a							)
          lowerCamelCase_						:   str		= output.audios[0]

          assert audio.ndim == 1
          assert len(__a							) == 256

          lowerCamelCase_						:   int		= audio[:10]
          lowerCamelCase_						:   Dict		= np.array(
              [-0.0_051, 0.0_050, -0.0_060, 0.0_034, -0.0_026, 0.0_033, -0.0_027, 0.0_033, -0.0_028, 0.0_032]							)

          assert np.abs(audio_slice - expected_slice							).max() < 1e-2



       def 				_lowerCAmelCase							(						self				:				Union[str, Any]							)					->Dict:
          lowerCamelCase_						:   Dict		= """cpu"""  # ensure determinism for the device-dependent torch.Generator
          lowerCamelCase_						:   List[Any]		= self.get_dummy_components()
          lowerCamelCase_						:   List[str]		= PNDMScheduler(skip_prk_steps=__a							)
          lowerCamelCase_						:   int		= AudioLDMPipeline(**__a							)
          lowerCamelCase_						:   Optional[int]		= audioldm_pipe.to(__a							)
          audioldm_pipe.set_progress_bar_config(disable=__a							)

          lowerCamelCase_						:   List[str]		= """A hammer hitting a wooden surface"""

          # test num_waveforms_per_prompt=1 (default)
          lowerCamelCase_						:   Union[str, Any]		= audioldm_pipe(__a		,				num_inference_steps=2							).audios

          assert audios.shape == (1, 256)

          # test num_waveforms_per_prompt=1 (default) for batch of prompts
          lowerCamelCase_						:   List[Any]		= 2
          lowerCamelCase_						:   Tuple		= audioldm_pipe([prompt] * batch_size		,				num_inference_steps=2							).audios

          assert audios.shape == (batch_size, 256)

          # test num_waveforms_per_prompt for single prompt
          lowerCamelCase_						:   Union[str, Any]		= 2
          lowerCamelCase_						:   List[str]		= audioldm_pipe(__a		,				num_inference_steps=2		,				num_waveforms_per_prompt=__a							).audios

          assert audios.shape == (num_waveforms_per_prompt, 256)

          # test num_waveforms_per_prompt for batch of prompts
          lowerCamelCase_						:   Any		= 2
          lowerCamelCase_						:   Optional[Any]		= audioldm_pipe(
              [prompt] * batch_size		,				num_inference_steps=2		,				num_waveforms_per_prompt=__a							).audios

          assert audios.shape == (batch_size * num_waveforms_per_prompt, 256)



       def 				_lowerCAmelCase							(						self				:				Tuple							)					->Optional[Any]:
          lowerCamelCase_						:   Optional[int]		= """cpu"""  # ensure determinism for the device-dependent torch.Generator
          lowerCamelCase_						:   Union[str, Any]		= self.get_dummy_components()
          lowerCamelCase_						:   Dict		= AudioLDMPipeline(**__a							)
          lowerCamelCase_						:   Optional[int]		= audioldm_pipe.to(__a							)
          audioldm_pipe.set_progress_bar_config(disable=__a							)
          lowerCamelCase_						:   str		= audioldm_pipe.vocoder.config.sampling_rate

          lowerCamelCase_						:   List[Any]		= self.get_dummy_inputs(__a							)
          lowerCamelCase_						:   int		= audioldm_pipe(audio_length_in_s=0.016		,				**__a							)
          lowerCamelCase_						:   Optional[int]		= output.audios[0]

          assert audio.ndim == 1
          assert len(__a							) / vocoder_sampling_rate == 0.016

          lowerCamelCase_						:   Dict		= audioldm_pipe(audio_length_in_s=0.032		,				**__a							)
          lowerCamelCase_						:   Any		= output.audios[0]

          assert audio.ndim == 1
          assert len(__a							) / vocoder_sampling_rate == 0.032



       def 				_lowerCAmelCase							(						self				:				Dict							)					->str:
          lowerCamelCase_						:   Dict		= self.get_dummy_components()
          lowerCamelCase_						:   List[Any]		= AudioLDMPipeline(**__a							)
          lowerCamelCase_						:   int		= audioldm_pipe.to(__a							)
          audioldm_pipe.set_progress_bar_config(disable=__a							)

          lowerCamelCase_						:   Any		= ["""hey"""]

          lowerCamelCase_						:   Any		= audioldm_pipe(__a		,				num_inference_steps=1							)
          lowerCamelCase_						:   List[Any]		= output.audios.shape
          assert audio_shape == (1, 256)

          lowerCamelCase_						:   str		= audioldm_pipe.vocoder.config
          config.model_in_dim *= 2
          lowerCamelCase_						:   List[Any]		= SpeechTaHifiGan(__a							).to(__a							)
          lowerCamelCase_						:   Dict		= audioldm_pipe(__a		,				num_inference_steps=1							)
          lowerCamelCase_						:   Tuple		= output.audios.shape
          # waveform shape is unchanged, we just have 2x the number of mel channels in the spectrogram
          assert audio_shape == (1, 256)



       def 				_lowerCAmelCase							(						self				:				Optional[int]							)					->Optional[Any]:
          self._test_attention_slicing_forward_pass(test_mean_pixel_difference=__a							)



       def 				_lowerCAmelCase							(						self				:				str							)					->Dict:
          self._test_inference_batch_single_identical(test_mean_pixel_difference=__a							)




       @unittest.skipIf(
           torch_device != """cuda""" or not is_xformers_available()		,				reason="""XFormers attention is only available with CUDA and `xformers` installed"""		,				)
       def 				_lowerCAmelCase							(						self				:				Dict							)					->Dict:
          self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=__a							)







@slow
class   SCREAMING_SNAKE_CASE_			(unittest.TestCase						):




       '''simple docstring'''



       def 				_lowerCAmelCase							(						self				:				List[Any]							)					->Optional[Any]:
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()



       def 				_lowerCAmelCase							(						self				:				List[str]		,				__a				:				Dict		,				__a				:				Optional[Any]="cpu"		,				__a				:				Tuple=torch.floataa		,				__a				:				Optional[int]=0							)					->Dict:
          lowerCamelCase_						:   Union[str, Any]		= torch.Generator(device=__a							).manual_seed(__a							)
          lowerCamelCase_						:   int		= np.random.RandomState(__a							).standard_normal((1, 8, 128, 16)							)
          lowerCamelCase_						:   Tuple		= torch.from_numpy(__a							).to(device=__a		,				dtype=__a							)
          lowerCamelCase_						:   Tuple		= {
              """prompt""": """A hammer hitting a wooden surface""",
              """latents""": latents,
              """generator""": generator,
              """num_inference_steps""": 3,
              """guidance_scale""": 2.5,
          }
          return inputs



       def 				_lowerCAmelCase							(						self				:				Optional[int]							)					->Tuple:
          lowerCamelCase_						:   List[str]		= AudioLDMPipeline.from_pretrained("""cvssp/audioldm"""							)
          lowerCamelCase_						:   int		= audioldm_pipe.to(__a							)
          audioldm_pipe.set_progress_bar_config(disable=__a							)

          lowerCamelCase_						:   Optional[int]		= self.get_inputs(__a							)
          lowerCamelCase_						:   List[Any]		= 25
          lowerCamelCase_						:   Union[str, Any]		= audioldm_pipe(**__a							).audios[0]

          assert audio.ndim == 1
          assert len(__a							) == 81_920

          lowerCamelCase_						:   Tuple		= audio[77_230:77_240]
          lowerCamelCase_						:   Tuple		= np.array(
              [-0.4_884, -0.4_607, 0.0_023, 0.5_007, 0.5_896, 0.5_151, 0.3_813, -0.0_208, -0.3_687, -0.4_315]							)
          lowerCamelCase_						:   Dict		= np.abs(expected_slice - audio_slice							).max()
          assert max_diff < 1e-2




       def 				_lowerCAmelCase							(						self				:				Any							)					->str:
          lowerCamelCase_						:   List[Any]		= AudioLDMPipeline.from_pretrained("""cvssp/audioldm"""							)
          lowerCamelCase_						:   Dict		= LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config							)
          lowerCamelCase_						:   Union[str, Any]		= audioldm_pipe.to(__a							)
          audioldm_pipe.set_progress_bar_config(disable=__a							)

          lowerCamelCase_						:   Tuple		= self.get_inputs(__a							)
          lowerCamelCase_						:   List[str]		= audioldm_pipe(**__a							).audios[0]

          assert audio.ndim == 1
          assert len(__a							) == 81_920

          lowerCamelCase_						:   int		= audio[27_780:27_790]
          lowerCamelCase_						:   Optional[Any]		= np.array([-0.2_131, -0.0_873, -0.0_124, -0.0_189, 0.0_569, 0.1_373, 0.1_883, 0.2_886, 0.3_297, 0.2_212]							)
          lowerCamelCase_						:   Union[str, Any]		= np.abs(expected_slice - audio_slice							).max()
          assert max_diff < 3e-2









import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


snake_case__				:	Dict					       =					logging.getLogger(__name__)


def    __lowerCamelCase   (      )     ->    Any:
   lowerCamelCase_						:   str		= 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=1000						,						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=512						,						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."""						,						)

   lowerCamelCase_						:   Optional[int]		= parser.parse_args()
   return args


def    __lowerCamelCase   (      A__						: Tuple					)     ->    str:
   def fn(A__						: Dict					):
      return tokenizer(examples["""text"""]					)

   return fn


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


def    __lowerCamelCase   (      A__						: Tuple					)     ->    Union[str, Any]:
   lowerCamelCase_						:   int		= datasets.load_dataset(args.dataset_name						,						args.dataset_config						,						split=args.split					)

   if args.limit is not None:
      lowerCamelCase_						:   str		= min(len(A__					)						,						args.limit					)
      lowerCamelCase_						:   Dict		= dataset.select(range(A__					)					)
      print(f'''Limiting the dataset to {args.limit} entries.'''					)

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

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

   lowerCamelCase_						:   Optional[int]		= dataset_tokenized.map(A__						,						batched=A__						,						batch_size=1000						,						num_proc=4					)

   lowerCamelCase_						:   List[str]		= 0
   lowerCamelCase_						:   Tuple		= 0
   for shard in range(0						,						len(A__					)						,						args.shard_size					):
      lowerCamelCase_						:   Union[str, Any]		= grouped_dataset[shard : shard + args.shard_size]
      lowerCamelCase_						:   Optional[int]		= len(dataset_snapshot["""input_ids"""]					)
      lowerCamelCase_						:   List[str]		= os.path.join(A__						,						f'''dataset-{shard_count}-{records_containing}.tfrecord'''					)
      lowerCamelCase_						:   Optional[int]		= get_serialized_examples(A__					)

      with tf.io.TFRecordWriter(A__					) as out_file:
         for i in range(len(A__					)					):
            lowerCamelCase_						:   Dict		= 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__":
      snake_case__				:	str					       =					parse_args()
      main(args)





