






"""simple docstring"""






import math
import time
from typing import Dict, List, Optional

from torch.utils.data import Dataset

from transformers import SeqaSeqTrainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput, speed_metrics


if is_torch_tpu_available(check_device=False):
  import torch_xla.core.xla_model as xm
  import torch_xla.debug.metrics as met




class 							lowerCAmelCase__       (    UpperCAmelCase__				):






      '''simple docstring'''

      def __init__(		self  :						Tuple		,			*lowercase_  :						Any		,			lowercase_  :						Dict=None		,			lowercase_  :						Dict=None		,			**lowercase_  :						List[str]):




          '''simple docstring'''


          super().__init__(*lowercase_		,			**lowercase_)
          SCREAMING_SNAKE_CASE_   :						Dict							=	eval_examples
          SCREAMING_SNAKE_CASE_   :						Any							=	post_process_function
      def _SCREAMING_SNAKE_CASE				(		self  :						Dict		,			lowercase_  :						Optional[Dataset] = None		,			lowercase_  :						List[Any]=None		,			lowercase_  :						Optional[List[str]] = None		,			lowercase_  :						str = "eval"		,			**lowercase_  :						List[Any]		,			):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	gen_kwargs.copy()
          SCREAMING_SNAKE_CASE_   :						int							=	(
              gen_kwargs['''max_length'''] if gen_kwargs.get('''max_length''') is not None else self.args.generation_max_length
          )
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	(
              gen_kwargs['''num_beams'''] if gen_kwargs.get('''num_beams''') is not None else self.args.generation_num_beams
          )
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	gen_kwargs

          SCREAMING_SNAKE_CASE_   :						Tuple							=	self.eval_dataset if eval_dataset is None else eval_dataset
          SCREAMING_SNAKE_CASE_   :						Dict							=	self.get_eval_dataloader(lowercase_)
          SCREAMING_SNAKE_CASE_   :						List[str]							=	self.eval_examples if eval_examples is None else eval_examples

          # Temporarily disable metric computation, we will do it in the loop here.
          SCREAMING_SNAKE_CASE_   :						Any							=	self.compute_metrics
          SCREAMING_SNAKE_CASE_   :						int							=	None
          SCREAMING_SNAKE_CASE_   :						str							=	time.time()
          SCREAMING_SNAKE_CASE_   :						Any							=	self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
          try:
              SCREAMING_SNAKE_CASE_   :						List[str]							=	eval_loop(
                  lowercase_		,			description='''Evaluation'''		,			prediction_loss_only=True if compute_metrics is None else None		,			ignore_keys=lowercase_		,			metric_key_prefix=lowercase_		,			)
          finally:
              SCREAMING_SNAKE_CASE_   :						List[Any]							=	compute_metrics
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	self.args.eval_batch_size * self.args.world_size
          if F'{metric_key_prefix}_jit_compilation_time' in output.metrics:
              start_time += output.metrics[F'{metric_key_prefix}_jit_compilation_time']
          output.metrics.update(
              speed_metrics(
                  lowercase_		,			lowercase_		,			num_samples=output.num_samples		,			num_steps=math.ceil(output.num_samples / total_batch_size)		,			))

          if self.post_process_function is not None and self.compute_metrics is not None and self.args.should_save:
              # Only the main node write the results by default
              SCREAMING_SNAKE_CASE_   :						List[Any]							=	self.post_process_function(lowercase_		,			lowercase_		,			lowercase_)
              SCREAMING_SNAKE_CASE_   :						List[Any]							=	self.compute_metrics(lowercase_)

              # Prefix all keys with metric_key_prefix + '_'
              for key in list(metrics.keys()):
                  if not key.startswith(F'{metric_key_prefix}_'):
                      SCREAMING_SNAKE_CASE_   :						Optional[int]							=	metrics.pop(lowercase_)

              metrics.update(output.metrics)
          else:
              SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	output.metrics

          if self.args.should_log:
              # Only the main node log the results by default
              self.log(lowercase_)

          if self.args.tpu_metrics_debug or self.args.debug:
              # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
              xm.master_print(met.metrics_report())

          SCREAMING_SNAKE_CASE_   :						Tuple							=	self.callback_handler.on_evaluate(self.args		,			self.state		,			self.control		,			lowercase_)
          return metrics



      def _SCREAMING_SNAKE_CASE				(		self  :						Optional[int]		,			lowercase_  :						int		,			lowercase_  :						Dict		,			lowercase_  :						Union[str, Any]=None		,			lowercase_  :						str = "test"		,			**lowercase_  :						List[str]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	gen_kwargs.copy()

          SCREAMING_SNAKE_CASE_   :						List[Any]							=	self.get_test_dataloader(lowercase_)

          # Temporarily disable metric computation, we will do it in the loop here.
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	self.compute_metrics
          SCREAMING_SNAKE_CASE_   :						str							=	None
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	time.time()
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
          try:
              SCREAMING_SNAKE_CASE_   :						List[Any]							=	eval_loop(
                  lowercase_		,			description='''Prediction'''		,			prediction_loss_only=True if compute_metrics is None else None		,			ignore_keys=lowercase_		,			metric_key_prefix=lowercase_		,			)
          finally:
              SCREAMING_SNAKE_CASE_   :						Optional[int]							=	compute_metrics

          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	self.args.eval_batch_size * self.args.world_size
          if F'{metric_key_prefix}_jit_compilation_time' in output.metrics:
              start_time += output.metrics[F'{metric_key_prefix}_jit_compilation_time']
          output.metrics.update(
              speed_metrics(
                  lowercase_		,			lowercase_		,			num_samples=output.num_samples		,			num_steps=math.ceil(output.num_samples / total_batch_size)		,			))
          if self.post_process_function is None or self.compute_metrics is None:
              return output

          SCREAMING_SNAKE_CASE_   :						Any							=	self.post_process_function(lowercase_		,			lowercase_		,			lowercase_		,			'''predict''')
          SCREAMING_SNAKE_CASE_   :						str							=	self.compute_metrics(lowercase_)

          # Prefix all keys with metric_key_prefix + '_'
          for key in list(metrics.keys()):
              if not key.startswith(F'{metric_key_prefix}_'):
                  SCREAMING_SNAKE_CASE_   :						Any							=	metrics.pop(lowercase_)
          metrics.update(output.metrics)
          return PredictionOutput(predictions=predictions.predictions		,			label_ids=predictions.label_ids		,			metrics=lowercase_)










"""simple docstring"""






import multiprocessing
import time

from arguments import PretokenizationArguments
from datasets import load_dataset

from transformers import AutoTokenizer, HfArgumentParser



def       _A    (__a							)				-> Optional[int]:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	{}
    SCREAMING_SNAKE_CASE_   :						List[Any]							=	tokenizer(example['''content''']					,       truncation=__a							)['''input_ids''']
    SCREAMING_SNAKE_CASE_   :						List[Any]							=	len(example['''content''']							) / len(output['''input_ids''']							)
    return output


UpperCAmelCase_     :				Tuple         =				HfArgumentParser(PretokenizationArguments)
UpperCAmelCase_     :				Tuple         =				parser.parse_args()
if args.num_workers is None:
  UpperCAmelCase_     :				Tuple         =				multiprocessing.cpu_count()
UpperCAmelCase_     :				str         =				AutoTokenizer.from_pretrained(args.tokenizer_dir)

UpperCAmelCase_     :				Tuple         =				time.time()
UpperCAmelCase_     :				List[Any]         =				load_dataset(args.dataset_name, split="""train""")
print(f'''Dataset loaded in {time.time()-t_start:.2f}s''')

UpperCAmelCase_     :				Dict         =				time.time()
UpperCAmelCase_     :				Dict         =				ds.map(
    tokenize,
    num_proc=args.num_workers,
    remove_columns=[
        """repo_name""",
        """path""",
        """copies""",
        """size""",
        """content""",
        """license""",
        """hash""",
        """line_mean""",
        """line_max""",
        """alpha_frac""",
        """autogenerated""",
    ],
)
print(f'''Dataset tokenized in {time.time()-t_start:.2f}s''')

UpperCAmelCase_     :				str         =				time.time()
ds.push_to_hub(args.tokenized_data_repo)





print(f'''Data pushed to the hub in {time.time()-t_start:.2f}s''')


