





'''simple docstring'''






import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


__lowerCamelCase						: int		=   logging.getLogger(__name__)


def 					_a   ():


  """simple docstring"""

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


  """simple docstring"""

  def fn(__SCREAMING_SNAKE_CASE   ):
    return tokenizer(examples['''text''']   )

  return fn


def 					_a   (__SCREAMING_SNAKE_CASE   ):


  """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=__SCREAMING_SNAKE_CASE   )
    _UpperCamelCase       			=tf.train.Example(features=__SCREAMING_SNAKE_CASE   )
    _UpperCamelCase       			=example.SerializeToString()
    records.append(__SCREAMING_SNAKE_CASE   )
  return records


def 					_a   (__SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

  _UpperCamelCase       			=datasets.load_dataset(args.dataset_name  ,      args.dataset_config  ,      split=args.split   )

  if args.limit is not None:
    _UpperCamelCase       			=min(len(__SCREAMING_SNAKE_CASE   )  ,      args.limit   )
    _UpperCamelCase       			=dataset.select(range(__SCREAMING_SNAKE_CASE   )   )
    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(__SCREAMING_SNAKE_CASE   ):
      os.makedirs(__SCREAMING_SNAKE_CASE   )
  else:
    _UpperCamelCase       			=os.path.join(args.output_dir  ,      args.split   )

  # Tokenize the whole dataset at once.
  _UpperCamelCase       			=tokenize_function(__SCREAMING_SNAKE_CASE   )
  _UpperCamelCase       			=dataset.map(__SCREAMING_SNAKE_CASE  ,      batched=__SCREAMING_SNAKE_CASE  ,      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(__SCREAMING_SNAKE_CASE   ):
    # 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  ,      __SCREAMING_SNAKE_CASE  ,      args.max_length   )]
        for k, t in concatenated_examples.items()
    }
    return result

  _UpperCamelCase       			=dataset_tokenized.map(__SCREAMING_SNAKE_CASE  ,      batched=__SCREAMING_SNAKE_CASE  ,      batch_size=1000  ,      num_proc=4   )

  _UpperCamelCase       			=0
  _UpperCamelCase       			=0
  for shard in range(0  ,      len(__SCREAMING_SNAKE_CASE   )  ,      args.shard_size   ):
    _UpperCamelCase       			=grouped_dataset[shard : shard + args.shard_size]
    _UpperCamelCase       			=len(dataset_snapshot['''input_ids''']   )
    _UpperCamelCase       			=os.path.join(__SCREAMING_SNAKE_CASE  ,      f'''dataset-{shard_count}-{records_containing}.tfrecord'''   )
    _UpperCamelCase       			=get_serialized_examples(__SCREAMING_SNAKE_CASE   )

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

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


if __name__ == "__main__":
     __lowerCamelCase						: List[str]		=   parse_args()
     main(args)













'''simple docstring'''






import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
      import torch

if is_vision_available():
      from PIL import Image

      from transformers import PoolFormerImageProcessor







class 		UpperCAmelCase   (				unittest.TestCase):


  """simple docstring"""







  def __init__(				self    :					Optional[int]							,	UpperCamelCase__    :					Union[str, Any]							,	UpperCamelCase__    :					Any=7							,	UpperCamelCase__    :					List[Any]=3							,	UpperCamelCase__    :					Dict=30							,	UpperCamelCase__    :					Union[str, Any]=400							,	UpperCamelCase__    :					List[Any]=True							,	UpperCamelCase__    :					List[str]=None							,	UpperCamelCase__    :					List[str]=0.9							,	UpperCamelCase__    :					Any=None							,	UpperCamelCase__    :					Any=True							,	UpperCamelCase__    :					Optional[Any]=[0.5, 0.5, 0.5]							,	UpperCamelCase__    :					List[str]=[0.5, 0.5, 0.5]							,	)     ->						Dict:
    _UpperCamelCase       			=size if size is not None else {'''shortest_edge''': 30}
    _UpperCamelCase       			=crop_size if crop_size is not None else {'''height''': 30, '''width''': 30}
    _UpperCamelCase       			=parent
    _UpperCamelCase       			=batch_size
    _UpperCamelCase       			=num_channels
    _UpperCamelCase       			=min_resolution
    _UpperCamelCase       			=max_resolution
    _UpperCamelCase       			=do_resize_and_center_crop
    _UpperCamelCase       			=size
    _UpperCamelCase       			=crop_pct
    _UpperCamelCase       			=crop_size
    _UpperCamelCase       			=do_normalize
    _UpperCamelCase       			=image_mean
    _UpperCamelCase       			=image_std




  def   UpperCamelCase__      (				self    :					str  )     ->						Union[str, Any]:
    return {
        "size": self.size,
        "do_resize_and_center_crop": self.do_resize_and_center_crop,
        "crop_pct": self.crop_pct,
        "crop_size": self.crop_size,
        "do_normalize": self.do_normalize,
        "image_mean": self.image_mean,
        "image_std": self.image_std,
    }







@require_torch
@require_vision
class 		UpperCAmelCase   (				lowercase_   ,     unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				PoolFormerImageProcessor if is_vision_available() else None


  def   UpperCamelCase__      (				self    :					Optional[int]  )     ->						int:
    _UpperCamelCase       			=PoolFormerImageProcessingTester(self  )


  @property
  def   UpperCamelCase__      (				self    :					Tuple  )     ->						Optional[int]:
    return self.image_processor_tester.prepare_image_processor_dict()


  def   UpperCamelCase__      (				self    :					Optional[Any]  )     ->						int:
    _UpperCamelCase       			=self.image_processing_class(**self.image_processor_dict  )
    self.assertTrue(hasattr(UpperCamelCase__							,	'''do_resize_and_center_crop'''  )  )
    self.assertTrue(hasattr(UpperCamelCase__							,	'''size'''  )  )
    self.assertTrue(hasattr(UpperCamelCase__							,	'''crop_pct'''  )  )
    self.assertTrue(hasattr(UpperCamelCase__							,	'''do_normalize'''  )  )
    self.assertTrue(hasattr(UpperCamelCase__							,	'''image_mean'''  )  )
    self.assertTrue(hasattr(UpperCamelCase__							,	'''image_std'''  )  )


  def   UpperCamelCase__      (				self    :					Any  )     ->						Dict:
    _UpperCamelCase       			=self.image_processing_class.from_dict(self.image_processor_dict  )
    self.assertEqual(image_processor.size							,	{'''shortest_edge''': 30}  )
    self.assertEqual(image_processor.crop_size							,	{'''height''': 30, '''width''': 30}  )

    _UpperCamelCase       			=self.image_processing_class.from_dict(self.image_processor_dict							,	size=42							,	crop_size=84  )
    self.assertEqual(image_processor.size							,	{'''shortest_edge''': 42}  )
    self.assertEqual(image_processor.crop_size							,	{'''height''': 84, '''width''': 84}  )


  def   UpperCamelCase__      (				self    :					Dict  )     ->						Optional[Any]:
    pass


  def   UpperCamelCase__      (				self    :					Union[str, Any]  )     ->						List[Any]:
    # Initialize image_processing
    _UpperCamelCase       			=self.image_processing_class(**self.image_processor_dict  )
    # create random PIL images
    _UpperCamelCase       			=prepare_image_inputs(self.image_processor_tester							,	equal_resolution=UpperCamelCase__  )
    for image in image_inputs:
      self.assertIsInstance(UpperCamelCase__							,	Image.Image  )

    # Test not batched input
    _UpperCamelCase       			=image_processing(image_inputs[0]							,	return_tensors='''pt'''  ).pixel_values

    self.assertEqual(
        encoded_images.shape							,	(
            1,
            self.image_processor_tester.num_channels,
            self.image_processor_tester.crop_size['''height'''],
            self.image_processor_tester.crop_size['''width'''],
        )							,	)

    # Test batched
    _UpperCamelCase       			=image_processing(UpperCamelCase__							,	return_tensors='''pt'''  ).pixel_values
    self.assertEqual(
        encoded_images.shape							,	(
            self.image_processor_tester.batch_size,
            self.image_processor_tester.num_channels,
            self.image_processor_tester.crop_size['''height'''],
            self.image_processor_tester.crop_size['''width'''],
        )							,	)


  def   UpperCamelCase__      (				self    :					Union[str, Any]  )     ->						Optional[int]:
    # Initialize image_processing
    _UpperCamelCase       			=self.image_processing_class(**self.image_processor_dict  )
    # create random numpy tensors
    _UpperCamelCase       			=prepare_image_inputs(self.image_processor_tester							,	equal_resolution=UpperCamelCase__							,	numpify=UpperCamelCase__  )
    for image in image_inputs:
      self.assertIsInstance(UpperCamelCase__							,	np.ndarray  )

    # Test not batched input
    _UpperCamelCase       			=image_processing(image_inputs[0]							,	return_tensors='''pt'''  ).pixel_values
    self.assertEqual(
        encoded_images.shape							,	(
            1,
            self.image_processor_tester.num_channels,
            self.image_processor_tester.crop_size['''height'''],
            self.image_processor_tester.crop_size['''width'''],
        )							,	)

    # Test batched
    _UpperCamelCase       			=image_processing(UpperCamelCase__							,	return_tensors='''pt'''  ).pixel_values
    self.assertEqual(
        encoded_images.shape							,	(
            self.image_processor_tester.batch_size,
            self.image_processor_tester.num_channels,
            self.image_processor_tester.crop_size['''height'''],
            self.image_processor_tester.crop_size['''width'''],
        )							,	)




  def   UpperCamelCase__      (				self    :					str  )     ->						List[Any]:
    # Initialize image_processing
    _UpperCamelCase       			=self.image_processing_class(**self.image_processor_dict  )
    # create random PyTorch tensors
    _UpperCamelCase       			=prepare_image_inputs(self.image_processor_tester							,	equal_resolution=UpperCamelCase__							,	torchify=UpperCamelCase__  )
    for image in image_inputs:
      self.assertIsInstance(UpperCamelCase__							,	torch.Tensor  )

    # Test not batched input
    _UpperCamelCase       			=image_processing(image_inputs[0]							,	return_tensors='''pt'''  ).pixel_values
    self.assertEqual(
        encoded_images.shape							,	(
            1,
            self.image_processor_tester.num_channels,
            self.image_processor_tester.crop_size['''height'''],
            self.image_processor_tester.crop_size['''width'''],
        )							,	)

    # Test batched
    _UpperCamelCase       			=image_processing(UpperCamelCase__							,	return_tensors='''pt'''  ).pixel_values
    self.assertEqual(
        encoded_images.shape							,	(
            self.image_processor_tester.batch_size,
            self.image_processor_tester.num_channels,
            self.image_processor_tester.crop_size['''height'''],
            self.image_processor_tester.crop_size['''width'''],
        )							,	)






