


'''simple docstring'''





import gzip
import hashlib
import json
import multiprocessing
import os
import re
import shutil
import time
from pathlib import Path

import numpy as np
from arguments import PreprocessingArguments
from datasets import load_dataset
from minhash_deduplication import deduplicate_dataset

from transformers import AutoTokenizer, HfArgumentParser


__lowerCamelCase												=						re.compile(R'''\s+''')


def        a__     (					UpperCamelCase_						:      Any						):
   return {"hash": hashlib.mda(re.sub(UpperCamelCase_,   '''''',   example['''content''']						).encode('''utf-8'''						)						).hexdigest()}


def        a__     (					UpperCamelCase_						:      Any						):
   UpperCAmelCase__	:Any      =      [len(UpperCamelCase_						) for line in example['''content'''].splitlines()]
   return {"line_mean": np.mean(UpperCamelCase_						), "line_max": max(UpperCamelCase_						)}


def        a__     (					UpperCamelCase_						:      List[str]						):
   UpperCAmelCase__	:str      =      np.mean([c.isalnum() for c in example['''content''']]						)
   return {"alpha_frac": alpha_frac}


def        a__     (					UpperCamelCase_						:      Any,   UpperCamelCase_						:      Any						):
   if example["hash"] in uniques:
      uniques.remove(example['''hash''']						)
      return True
   else:
      return False


def        a__     (					UpperCamelCase_						:      Tuple,   UpperCamelCase_						:      Tuple=5						):
   UpperCAmelCase__	:Optional[int]      =      ['''auto-generated''', '''autogenerated''', '''automatically generated''']
   UpperCAmelCase__	:int      =      example['''content'''].splitlines()
   for _, line in zip(range(UpperCamelCase_						),   UpperCamelCase_						):
      for keyword in keywords:
         if keyword in line.lower():
            return {"autogenerated": True}
   else:
      return {"autogenerated": False}


def        a__     (					UpperCamelCase_						:      Dict,   UpperCamelCase_						:      Any=5,   UpperCamelCase_						:      Dict=0.05						):
   UpperCAmelCase__	:List[Any]      =      ['''unit tests''', '''test file''', '''configuration file''']
   UpperCAmelCase__	:Tuple      =      example['''content'''].splitlines()
   UpperCAmelCase__	:List[Any]      =      0
   UpperCAmelCase__	:str      =      0
   # first test
   for _, line in zip(range(UpperCamelCase_						),   UpperCamelCase_						):
      for keyword in keywords:
         if keyword in line.lower():
            return {"config_or_test": True}
    # second test
   UpperCAmelCase__	:Optional[int]      =      example['''content'''].count('''\n'''						)
   UpperCAmelCase__	:List[str]      =      int(coeff * nlines						)
   for line in lines:
      count_config += line.lower().count('''config'''						)
      count_test += line.lower().count('''test'''						)
      if count_config > threshold or count_test > threshold:
         return {"config_or_test": True}
   return {"config_or_test": False}


def        a__     (					UpperCamelCase_						:      Optional[int]						):
   UpperCAmelCase__	:Optional[int]      =      ['''def ''', '''class ''', '''for ''', '''while ''']
   UpperCAmelCase__	:str      =      example['''content'''].splitlines()
   for line in lines:
      for keyword in keywords:
         if keyword in line.lower():
            return {"has_no_keywords": False}
   return {"has_no_keywords": True}


def        a__     (					UpperCamelCase_						:      int,   UpperCamelCase_						:      Dict=4						):
   UpperCAmelCase__	:Dict      =      example['''content'''].splitlines()
   UpperCAmelCase__	:Union[str, Any]      =      0
   for line in lines:
      counter += line.lower().count('''='''						)
      if counter > minimum:
         return {"has_few_assignments": False}
   return {"has_few_assignments": True}


def        a__     (					UpperCamelCase_						:      Tuple						):
   UpperCAmelCase__	:List[Any]      =      tokenizer(example['''content'''],   truncation=UpperCamelCase_						)['''input_ids''']
   UpperCAmelCase__	:int      =      len(example['''content''']						) / len(UpperCamelCase_						)
   return {"ratio": ratio}


def        a__     (					UpperCamelCase_						:      str						):
   UpperCAmelCase__	:List[Any]      =      {}
   results.update(get_hash(UpperCamelCase_						)						)
   results.update(line_stats(UpperCamelCase_						)						)
   results.update(alpha_stats(UpperCamelCase_						)						)
   results.update(char_token_ratio(UpperCamelCase_						)						)
   results.update(is_autogenerated(UpperCamelCase_						)						)
   results.update(is_config_or_test(UpperCamelCase_						)						)
   results.update(has_no_keywords(UpperCamelCase_						)						)
   results.update(has_few_assignments(UpperCamelCase_						)						)
   return results


def        a__     (					UpperCamelCase_						:      Dict,   UpperCamelCase_						:      List[Any],   UpperCamelCase_						:      List[str]						):
   if not check_uniques(UpperCamelCase_,   UpperCamelCase_						):
      return False
   elif example["autogenerated"]:
      return False
   elif example["line_max"] > args.line_max:
      return False
   elif example["line_mean"] > args.line_mean:
      return False
   elif example["alpha_frac"] < args.alpha_frac:
      return False
   elif example["ratio"] < args.min_token_ratio:
      return False
   elif example["config_or_test"] and np.random.rand() <= args.filter_proba:
      return False
   elif example["has_no_keywords"] and np.random.rand() <= args.filter_proba:
      return False
   elif example["has_few_assignments"]:
      return False
   else:
      return True


def        a__     (					UpperCamelCase_						:      Optional[Any]						):
   with open(UpperCamelCase_,   '''rb'''						) as f_in:
      with gzip.open(str(UpperCamelCase_						) + '''.gz''',   '''wb''',   compresslevel=6						) as f_out:
         shutil.copyfileobj(UpperCamelCase_,   UpperCamelCase_						)
   os.unlink(UpperCamelCase_						)


# Settings
__lowerCamelCase												=						HfArgumentParser(PreprocessingArguments)
__lowerCamelCase												=						parser.parse_args()
if args.num_workers is None:
    __lowerCamelCase												=						multiprocessing.cpu_count()
__lowerCamelCase												=						AutoTokenizer.from_pretrained(args.tokenizer_dir)

# Load dataset
__lowerCamelCase												=						time.time()
__lowerCamelCase												=						load_dataset(args.dataset_name, split='''train''')
print(F'''Time to load dataset: {time.time()-t_start:.2f}''')

# Run preprocessing
__lowerCamelCase												=						time.time()
__lowerCamelCase												=						ds.map(preprocess, num_proc=args.num_workers)
print(F'''Time to preprocess dataset: {time.time()-t_start:.2f}''')

# Deduplicate hashes
__lowerCamelCase												=						set(ds.unique('''hash'''))
__lowerCamelCase												=						len(uniques) / len(ds)
print(F'''Fraction of duplicates: {1-frac:.2%}''')

# Deduplicate data and apply heuristics
__lowerCamelCase												=						time.time()
__lowerCamelCase												=						ds.filter(filter, fn_kwargs={'''uniques''': uniques, '''args''': args})
print(F'''Time to filter dataset: {time.time()-t_start:.2f}''')
print(F'''Size of filtered dataset: {len(ds_filter)}''')

# Deduplicate with minhash and jaccard similarity
if args.near_deduplication:
    __lowerCamelCase												=						time.time()
    __lowerCamelCase					,		__lowerCamelCase												=						deduplicate_dataset(ds_filter, args.jaccard_threshold)
    print(F'''Time to deduplicate dataset: {time.time()-t_start:.2f}''')
    print(F'''Size of deduplicate dataset: {len(ds_filter)}''')

# Save data in batches of samples_per_file
__lowerCamelCase												=						Path(args.output_dir)
output_dir.mkdir(exist_ok=True)

# save duplicate_clusters in the output_dir as artifacts
# not sure it is the right place the save it
if args.near_deduplication:
    with open(output_dir / '''duplicate_clusters.json''', '''w''') as f:
        json.dump(duplicate_clusters, f)

__lowerCamelCase												=						output_dir / '''data'''
data_dir.mkdir(exist_ok=True)

__lowerCamelCase												=						time.time()
for file_number, index in enumerate(range(0, len(ds_filter), args.samples_per_file)):
    __lowerCamelCase												=						str(data_dir / F'''file-{file_number+1:012}.json''')
    __lowerCamelCase												=						min(len(ds_filter), index + args.samples_per_file)
    ds_filter.select(list(range(index, end_index))).to_json(file_path)
    compress_file(file_path)


print(F'''Time to save dataset: {time.time()-t_start:.2f}''')







'''simple docstring'''





import importlib.util
import json
import os
import warnings
from dataclasses import dataclass, field

import torch

from ..training_args import TrainingArguments
from ..utils import cached_property, is_sagemaker_dp_enabled, logging


__lowerCamelCase												=						logging.get_logger(__name__)


def        a__     (					):
   # Get the sagemaker specific mp parameters from smp_options variable.
   UpperCAmelCase__	:Any      =      os.getenv('''SM_HP_MP_PARAMETERS''',   '''{}'''						)
   try:
      # Parse it and check the field "partitions" is included, it is required for model parallel.
      UpperCAmelCase__	:Any      =      json.loads(UpperCamelCase_						)
      if "partitions" not in smp_options:
         return False
   except json.JSONDecodeError:
      return False

   # Get the sagemaker specific framework parameters from mpi_options variable.
   UpperCAmelCase__	:Optional[Any]      =      os.getenv('''SM_FRAMEWORK_PARAMS''',   '''{}'''						)
   try:
      # Parse it and check the field "sagemaker_distributed_dataparallel_enabled".
      UpperCAmelCase__	:str      =      json.loads(UpperCamelCase_						)
      if not mpi_options.get('''sagemaker_mpi_enabled''',   UpperCamelCase_						):
         return False
   except json.JSONDecodeError:
      return False
   # Lastly, check if the `smdistributed` module is present.
   return importlib.util.find_spec('''smdistributed'''						) is not None


if is_sagemaker_model_parallel_available():
    import smdistributed.modelparallel.torch as smp

    smp.init()




@dataclass
class 				UpperCAmelCase			(			_snake_case						):
 UpperCAmelCase				=					field(
     default=""						, metadata={"help": "Used by the SageMaker launcher to send mp-specific args. Ignored in SageMakerTrainer"}						, )


 def        __SCREAMING_SNAKE_CASE					(      self							:	List[Any] ):
    super().__post_init__()
    warnings.warn(
        '''`SageMakerTrainingArguments` is deprecated and will be removed in v5 of Transformers. You can use '''
        '''`TrainingArguments` instead.'''			,       __lowerCamelCase			,       )


 @cached_property
 def        __SCREAMING_SNAKE_CASE					(      self							:	List[Any] ):
    logger.info('''PyTorch: setting up devices''' )
    if torch.distributed.is_available() and torch.distributed.is_initialized() and self.local_rank == -1:
       logger.warning(
           '''torch.distributed process group is initialized, but local_rank == -1. '''
           '''In order to use Torch DDP, launch your script with `python -m torch.distributed.launch''' )
    if self.no_cuda:
       UpperCAmelCase__	:Union[str, Any]      =      torch.device('''cpu''' )
       UpperCAmelCase__	:Dict      =      0
    elif is_sagemaker_model_parallel_available():
       UpperCAmelCase__	:Dict      =      smp.local_rank()
       UpperCAmelCase__	:List[str]      =      torch.device('''cuda'''			,       __lowerCamelCase )
       UpperCAmelCase__	:Optional[int]      =      1
    elif is_sagemaker_dp_enabled():
       import smdistributed.dataparallel.torch.torch_smddp  # noqa: F401

       torch.distributed.init_process_group(backend='''smddp'''			,       timeout=self.ddp_timeout_delta )
       UpperCAmelCase__	:int      =      int(os.getenv('''SMDATAPARALLEL_LOCAL_RANK''' ) )
       UpperCAmelCase__	:int      =      torch.device('''cuda'''			,       self.local_rank )
       UpperCAmelCase__	:Union[str, Any]      =      1
    elif self.local_rank == -1:
       # if n_gpu is > 1 we'll use nn.DataParallel.
       # If you only want to use a specific subset of GPUs use `CUDA_VISIBLE_DEVICES=0`
       # Explicitly set CUDA to the first (index 0) CUDA device, otherwise `set_device` will
       # trigger an error that a device index is missing. Index 0 takes into account the
       # GPUs available in the environment, so `CUDA_VISIBLE_DEVICES=1,2` with `cuda:0`
       # will use the first GPU in that env, i.e. GPU#1
       UpperCAmelCase__	:str      =      torch.device('''cuda:0''' if torch.cuda.is_available() else '''cpu''' )
       # Sometimes the line in the postinit has not been run before we end up here, so just checking we're not at
       # the default value.
       UpperCAmelCase__	:str      =      torch.cuda.device_count()
    else:
       # Here, we'll use torch.distributed.
       # Initializes the distributed backend which will take care of synchronizing nodes/GPUs
       if not torch.distributed.is_initialized():
          torch.distributed.init_process_group(backend='''nccl'''			,       timeout=self.ddp_timeout_delta )
       UpperCAmelCase__	:str      =      torch.device('''cuda'''			,       self.local_rank )
       UpperCAmelCase__	:Union[str, Any]      =      1

    if device.type == "cuda":
       torch.cuda.set_device(__lowerCamelCase )

    return device


 @property
 def        __SCREAMING_SNAKE_CASE					(      self							:	List[str] ):
    if is_sagemaker_model_parallel_available():
       return smp.dp_size()

    return super().world_size


 @property
 def        __SCREAMING_SNAKE_CASE					(      self							:	Dict ):
    return not is_sagemaker_model_parallel_available()





 @property
 def        __SCREAMING_SNAKE_CASE					(      self							:	int ):
    return False



