


"""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


__snake_case   			=    re.compile(R"""\s+""")







def   __lowerCAmelCase						(						lowercase  :  Optional[int] )	->						Optional[int]:


    """simple docstring"""





    return {"hash": hashlib.mda(re.sub(lowercase						,							""						,							example["content"] ).encode("utf-8" ) ).hexdigest()}







def   __lowerCAmelCase						(						lowercase  :  int )	->						List[Any]:


    """simple docstring"""





    snake_case    :      Optional[int]      					=    [len(lowercase ) for line in example["content"].splitlines()]
    return {"line_mean": np.mean(lowercase ), "line_max": max(lowercase )}







def   __lowerCAmelCase						(						lowercase  :  Optional[int] )	->						Any:


    """simple docstring"""





    snake_case    :      Union[str, Any]      					=    np.mean([c.isalnum() for c in example["content"]] )
    return {"alpha_frac": alpha_frac}







def   __lowerCAmelCase						(						lowercase  :  Optional[int]						,							lowercase  :  Any )	->						Union[str, Any]:


    """simple docstring"""





    if example["hash"] in uniques:
        uniques.remove(example["hash"] )
        return True
    else:
        return False







def   __lowerCAmelCase						(						lowercase  :  List[Any]						,							lowercase  :  Optional[int]=5 )	->						str:


    """simple docstring"""





    snake_case    :      List[str]      					=    ["auto-generated", "autogenerated", "automatically generated"]
    snake_case    :      List[str]      					=    example["content"].splitlines()
    for _, line in zip(range(lowercase )						,							lowercase ):
        for keyword in keywords:
            if keyword in line.lower():
                return {"autogenerated": True}
    else:
        return {"autogenerated": False}







def   __lowerCAmelCase						(						lowercase  :  Union[str, Any]						,							lowercase  :  str=5						,							lowercase  :  str=0.05 )	->						str:


    """simple docstring"""





    snake_case    :      str      					=    ["unit tests", "test file", "configuration file"]
    snake_case    :      Tuple      					=    example["content"].splitlines()
    snake_case    :      List[Any]      					=    0
    snake_case    :      Union[str, Any]      					=    0
    # first test
    for _, line in zip(range(lowercase )						,							lowercase ):
        for keyword in keywords:
            if keyword in line.lower():
                return {"config_or_test": True}
    # second test
    snake_case    :      List[str]      					=    example["content"].count("\n" )
    snake_case    :      Optional[Any]      					=    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   __lowerCAmelCase						(						lowercase  :  Dict )	->						Dict:


    """simple docstring"""





    snake_case    :      Union[str, Any]      					=    ["def ", "class ", "for ", "while "]
    snake_case    :      Dict      					=    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   __lowerCAmelCase						(						lowercase  :  List[Any]						,							lowercase  :  Union[str, Any]=4 )	->						int:


    """simple docstring"""





    snake_case    :      List[Any]      					=    example["content"].splitlines()
    snake_case    :      Optional[int]      					=    0
    for line in lines:
        counter += line.lower().count("=" )
        if counter > minimum:
            return {"has_few_assignments": False}
    return {"has_few_assignments": True}







def   __lowerCAmelCase						(						lowercase  :  str )	->						List[str]:


    """simple docstring"""





    snake_case    :      List[str]      					=    tokenizer(example["content"]						,							truncation=lowercase )["input_ids"]
    snake_case    :      Tuple      					=    len(example["content"] ) / len(lowercase )
    return {"ratio": ratio}







def   __lowerCAmelCase						(						lowercase  :  List[str] )	->						Dict:


    """simple docstring"""





    snake_case    :      Optional[Any]      					=    {}
    results.update(get_hash(lowercase ) )
    results.update(line_stats(lowercase ) )
    results.update(alpha_stats(lowercase ) )
    results.update(char_token_ratio(lowercase ) )
    results.update(is_autogenerated(lowercase ) )
    results.update(is_config_or_test(lowercase ) )
    results.update(has_no_keywords(lowercase ) )
    results.update(has_few_assignments(lowercase ) )
    return results







def   __lowerCAmelCase						(						lowercase  :  Optional[Any]						,							lowercase  :  Optional[Any]						,							lowercase  :  Optional[Any] )	->						List[Any]:


    """simple docstring"""





    if not check_uniques(lowercase						,							lowercase ):
        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   __lowerCAmelCase						(						lowercase  :  List[str] )	->						List[str]:


    """simple docstring"""





    with open(lowercase						,							"rb" ) as f_in:
        with gzip.open(str(lowercase ) + ".gz"						,							"wb"						,							compresslevel=6 ) as f_out:
            shutil.copyfileobj(lowercase						,							lowercase )
    os.unlink(lowercase )


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

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

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

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

# Deduplicate data and apply heuristics
__snake_case   			=    time.time()
__snake_case   			=    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:
  __snake_case   			=    time.time()
  __snake_case			,		__snake_case   			=    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
__snake_case   			=    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)

__snake_case   			=    output_dir / """data"""
data_dir.mkdir(exist_ok=True)

__snake_case   			=    time.time()
for file_number, index in enumerate(range(0, len(ds_filter), args.samples_per_file)):
  __snake_case   			=    str(data_dir / F'''file-{file_number+1:012}.json''')
  __snake_case   			=    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


__snake_case   			=    logging.get_logger(__name__)







def   __lowerCAmelCase						(						)	->						str:


    """simple docstring"""





    snake_case    :      Dict      					=    os.getenv("SM_HP_MP_PARAMETERS"						,							"{}" )
    try:
        # Parse it and check the field "partitions" is included, it is required for model parallel.
        snake_case    :      Optional[int]      					=    json.loads(lowercase )
        if "partitions" not in smp_options:
            return False
    except json.JSONDecodeError:
        return False

    # Get the sagemaker specific framework parameters from mpi_options variable.
    snake_case    :      Optional[int]      					=    os.getenv("SM_FRAMEWORK_PARAMS"						,							"{}" )
    try:
        # Parse it and check the field "sagemaker_distributed_dataparallel_enabled".
        snake_case    :      Any      					=    json.loads(lowercase )
        if not mpi_options.get("sagemaker_mpi_enabled"						,							lowercase ):
            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 _lowerCAmelCase				(      snake_case_	):
   __UpperCAmelCase :       str							  =  field(
       default=''''''    ,							metadata={'''help''': '''Used by the SageMaker launcher to send mp-specific args. Ignored in SageMakerTrainer'''}    ,							)







   def  lowerCamelCase		(	self						)	->		List[Any]:

       '''simple docstring'''




       super().__post_init__()
       warnings.warn(
           "`SageMakerTrainingArguments` is deprecated and will be removed in v5 of Transformers. You can use "
           "`TrainingArguments` instead."						,	UpperCamelCase__						,	)







   @cached_property
   def  lowerCamelCase		(	self						)	->		"torch.device":

       '''simple docstring'''




       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:
           snake_case    :      Optional[Any]      					=    torch.device("cpu"						)
           snake_case    :      List[Any]      					=    0
       elif is_sagemaker_model_parallel_available():
           snake_case    :      Tuple      					=    smp.local_rank()
           snake_case    :      int      					=    torch.device("cuda"						,	UpperCamelCase__						)
           snake_case    :      Dict      					=    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						)
           snake_case    :      Any      					=    int(os.getenv("SMDATAPARALLEL_LOCAL_RANK"						)						)
           snake_case    :      Optional[Any]      					=    torch.device("cuda"						,	self.local_rank						)
           snake_case    :      str      					=    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
           snake_case    :      List[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.
           snake_case    :      Optional[Any]      					=    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						)
           snake_case    :      Any      					=    torch.device("cuda"						,	self.local_rank						)
           snake_case    :      Dict      					=    1

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

       return device







   @property
   def  lowerCamelCase		(	self						)	->		Optional[int]:

       '''simple docstring'''




       if is_sagemaker_model_parallel_available():
           return smp.dp_size()

       return super().world_size







   @property
   def  lowerCamelCase		(	self						)	->		List[str]:

       '''simple docstring'''




       return not is_sagemaker_model_parallel_available()







   @property
   def  lowerCamelCase		(	self						)	->		str:

       '''simple docstring'''




       return False





