






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 :      List[str]										=				re.compile(R"""\s+""")







def 					SCREAMING_SNAKE_CASE    (			snake_case_ :							List[str]       ) ->    str:
	return {"hash": hashlib.mda(re.sub(snake_case_					,  ""					,  example["content"]       ).encode("utf-8"       )       ).hexdigest()}







def 					SCREAMING_SNAKE_CASE    (			snake_case_ :							List[Any]       ) ->    int:
	snake_case__			:				str       =  [len(snake_case_       ) for line in example["content"].splitlines()]
	return {"line_mean": np.mean(snake_case_       ), "line_max": max(snake_case_       )}







def 					SCREAMING_SNAKE_CASE    (			snake_case_ :							List[str]       ) ->    int:
	snake_case__			:				Tuple       =  np.mean([c.isalnum() for c in example["content"]]       )
	return {"alpha_frac": alpha_frac}







def 					SCREAMING_SNAKE_CASE    (			snake_case_ :							Optional[int]					,  snake_case_ :							Any       ) ->    Dict:
	if example["hash"] in uniques:
		uniques.remove(example["hash"]       )
		return True
	else:
		return False







def 					SCREAMING_SNAKE_CASE    (			snake_case_ :							List[Any]					,  snake_case_ :							Dict=5       ) ->    List[Any]:
	snake_case__			:				Dict       =  ["auto-generated", "autogenerated", "automatically generated"]
	snake_case__			:				str       =  example["content"].splitlines()
	for _, line in zip(range(snake_case_       )					,  snake_case_       ):
		for keyword in keywords:
			if keyword in line.lower():
				return {"autogenerated": True}
	else:
		return {"autogenerated": False}







def 					SCREAMING_SNAKE_CASE    (			snake_case_ :							Tuple					,  snake_case_ :							List[str]=5					,  snake_case_ :							Dict=0.05       ) ->    int:
	snake_case__			:				int       =  ["unit tests", "test file", "configuration file"]
	snake_case__			:				Dict       =  example["content"].splitlines()
	snake_case__			:				List[str]       =  0
	snake_case__			:				Any       =  0
	# first test
	for _, line in zip(range(snake_case_       )					,  snake_case_       ):
		for keyword in keywords:
			if keyword in line.lower():
				return {"config_or_test": True}
    # second test
	snake_case__			:				Optional[int]       =  example["content"].count("\n"       )
	snake_case__			:				int       =  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 					SCREAMING_SNAKE_CASE    (			snake_case_ :							Optional[int]       ) ->    Optional[Any]:
	snake_case__			:				Any       =  ["def ", "class ", "for ", "while "]
	snake_case__			:				Union[str, Any]       =  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 					SCREAMING_SNAKE_CASE    (			snake_case_ :							int					,  snake_case_ :							int=4       ) ->    List[Any]:
	snake_case__			:				Optional[int]       =  example["content"].splitlines()
	snake_case__			:				Any       =  0
	for line in lines:
		counter += line.lower().count("="       )
		if counter > minimum:
			return {"has_few_assignments": False}
	return {"has_few_assignments": True}







def 					SCREAMING_SNAKE_CASE    (			snake_case_ :							List[str]       ) ->    int:
	snake_case__			:				str       =  tokenizer(example["content"]					,  truncation=snake_case_       )["input_ids"]
	snake_case__			:				int       =  len(example["content"]       ) / len(snake_case_       )
	return {"ratio": ratio}







def 					SCREAMING_SNAKE_CASE    (			snake_case_ :							Dict       ) ->    Dict:
	snake_case__			:				Dict       =  {}
	results.update(get_hash(snake_case_       )       )
	results.update(line_stats(snake_case_       )       )
	results.update(alpha_stats(snake_case_       )       )
	results.update(char_token_ratio(snake_case_       )       )
	results.update(is_autogenerated(snake_case_       )       )
	results.update(is_config_or_test(snake_case_       )       )
	results.update(has_no_keywords(snake_case_       )       )
	results.update(has_few_assignments(snake_case_       )       )
	return results







def 					SCREAMING_SNAKE_CASE    (			snake_case_ :							List[Any]					,  snake_case_ :							Union[str, Any]					,  snake_case_ :							Dict       ) ->    List[Any]:
	if not check_uniques(snake_case_					,  snake_case_       ):
		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 					SCREAMING_SNAKE_CASE    (			snake_case_ :							Tuple       ) ->    Any:
	with open(snake_case_					,  "rb"       ) as f_in:
		with gzip.open(str(snake_case_       ) + ".gz"					,  "wb"					,  compresslevel=6       ) as f_out:
			shutil.copyfileobj(snake_case_					,  snake_case_       )
	os.unlink(snake_case_       )


# Settings
__lowerCamelCase :      Optional[Any]										=				HfArgumentParser(PreprocessingArguments)
__lowerCamelCase :      List[str]										=				parser.parse_args()
if args.num_workers is None:
	__lowerCamelCase :      Optional[int]										=				multiprocessing.cpu_count()
__lowerCamelCase :      Dict										=				AutoTokenizer.from_pretrained(args.tokenizer_dir)

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

# Run preprocessing
__lowerCamelCase :      Tuple										=				time.time()
__lowerCamelCase :      Optional[int]										=				ds.map(preprocess, num_proc=args.num_workers)
print(f"Time to preprocess dataset: {time.time()-t_start:.2f}")

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

# Deduplicate data and apply heuristics
__lowerCamelCase :      List[Any]										=				time.time()
__lowerCamelCase :      str										=				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 :      Any										=				time.time()
	__lowerCamelCase :      Optional[Any]										=				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 :      Union[str, Any]										=				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 :      Tuple										=				output_dir / """data"""
data_dir.mkdir(exist_ok=True)

__lowerCamelCase :      Optional[Any]										=				time.time()
for file_number, index in enumerate(range(0, len(ds_filter), args.samples_per_file)):
	__lowerCamelCase :      Optional[Any]										=				str(data_dir / f"file-{file_number+1:012}.json")
	__lowerCamelCase :      List[str]										=				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}")











import json
from typing import Dict, List, Optional, Tuple, Union

from tokenizers import pre_tokenizers, processors

from ...tokenization_utils_base import AddedToken, BatchEncoding, EncodedInput
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import PaddingStrategy, logging
from .tokenization_led import LEDTokenizer


__lowerCamelCase :      int										=				logging.get_logger(__name__)


__lowerCamelCase :      List[str]										=				{"""vocab_file""": """vocab.json""", """merges_file""": """merges.txt""", """tokenizer_file""": """tokenizer.json"""}

__lowerCamelCase :      Any										=				{
    """vocab_file""": {
        """allenai/led-base-16384""": """https://huggingface.co/allenai/led-base-16384/resolve/main/vocab.json""",
    },
    """merges_file""": {
        """allenai/led-base-16384""": """https://huggingface.co/allenai/led-base-16384/resolve/main/merges.txt""",
    },
    """tokenizer_file""": {
        """allenai/led-base-16384""": """https://huggingface.co/allenai/led-base-16384/resolve/main/tokenizer.json""",
    },
}

__lowerCamelCase :      Optional[int]										=				{
    """allenai/led-base-16384""": 1_6384,
}



class 							SCREAMING_SNAKE_CASE__						(    UpperCamelCase_						):





					"""simple docstring"""






					a_     =		VOCAB_FILES_NAMES
					a_     =		PRETRAINED_VOCAB_FILES_MAP
					a_     =		PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
					a_     =		LEDTokenizer
					a_     =		["input_ids", "attention_mask"]



					def __init__( self					:  int							,		__A					:  Optional[int]=None							,		__A					:  Any=None							,		__A					:  Any=None							,		__A					:  int="replace"							,		__A					:  Optional[Any]="<s>"							,		__A					:  List[Any]="</s>"							,		__A					:  Dict="</s>"							,		__A					:  Union[str, Any]="<s>"							,		__A					:  List[Any]="<unk>"							,		__A					:  Optional[Any]="<pad>"							,		__A					:  Optional[int]="<mask>"							,		__A					:  str=False							,		__A					:  Optional[Any]=True							,		**__A					:  List[Any]							,		):
						super().__init__(
						    __A							,		__A							,		tokenizer_file=__A							,		errors=__A							,		bos_token=__A							,		eos_token=__A							,		sep_token=__A							,		cls_token=__A							,		unk_token=__A							,		pad_token=__A							,		mask_token=__A							,		add_prefix_space=__A							,		trim_offsets=__A							,		**__A							,		)

						snake_case__			:				List[Any]       =  json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()	)
						if pre_tok_state.get("add_prefix_space"							,		__A	) != add_prefix_space:
							snake_case__			:				List[str]       =  getattr(__A							,		pre_tok_state.pop("type"	)	)
							snake_case__			:				Tuple       =  add_prefix_space
							snake_case__			:				Any       =  pre_tok_class(**__A	)

						snake_case__			:				List[Any]       =  add_prefix_space

						# the pre_tokenizer is already updated in the GPT2TokenizerFast `__init__`
						snake_case__			:				str       =  "post_processor"
						snake_case__			:				Union[str, Any]       =  getattr(self.backend_tokenizer							,		__A							,		__A	)
						if tokenizer_component_instance:
							snake_case__			:				Any       =  json.loads(tokenizer_component_instance.__getstate__()	)

							# The lists 'sep' and 'cls' must be cased in tuples for the object `post_processor_class`
							if "sep" in state:
								snake_case__			:				int       =  tuple(state["sep"]	)
							if "cls" in state:
								snake_case__			:				Optional[int]       =  tuple(state["cls"]	)

							snake_case__			:				Union[str, Any]       =  False

							if state.get("add_prefix_space"							,		__A	) != add_prefix_space:
								snake_case__			:				Optional[Any]       =  add_prefix_space
								snake_case__			:				List[str]       =  True

							if state.get("trim_offsets"							,		__A	) != trim_offsets:
								snake_case__			:				str       =  trim_offsets
								snake_case__			:				Tuple       =  True

							if changes_to_apply:
								snake_case__			:				Any       =  getattr(__A							,		state.pop("type"	)	)
								snake_case__			:				Optional[Any]       =  component_class(**__A	)
								setattr(self.backend_tokenizer							,		__A							,		__A	)



					@property
					# Copied from transformers.models.bart.tokenization_bart_fast.BartTokenizerFast.mask_token with BART->LED
					def 	_lowercase					( self					:  Union[str, Any]	):
						if self._mask_token is None:
							if self.verbose:
								logger.error("Using mask_token, but it is not set yet."	)
							return None
						return str(self._mask_token	)



					@mask_token.setter
					def 	_lowercase					( self					:  str							,		__A					:  int	):
						snake_case__			:				str       =  AddedToken(__A							,		lstrip=__A							,		rstrip=__A	) if isinstance(__A							,		__A	) else value
						snake_case__			:				Optional[int]       =  value



					def 	_lowercase					( self					:  List[str]							,		*__A					:  Union[str, Any]							,		**__A					:  int	):
						snake_case__			:				Dict       =  kwargs.get("is_split_into_words"							,		__A	)

						if is_split_into_words and not self.add_prefix_space:
							raise ValueError(
							    f'''You need to instantiate {self.__class__.__name__} with add_prefix_space=True '''
							    "to use it with pretokenized inputs."	)

						return super()._batch_encode_plus(*__A							,		**__A	)



					def 	_lowercase					( self					:  Union[str, Any]							,		*__A					:  Tuple							,		**__A					:  Union[str, Any]	):
						snake_case__			:				List[str]       =  kwargs.get("is_split_into_words"							,		__A	)

						if is_split_into_words and not self.add_prefix_space:
							raise ValueError(
							    f'''You need to instantiate {self.__class__.__name__} with add_prefix_space=True '''
							    "to use it with pretokenized inputs."	)

						return super()._encode_plus(*__A							,		**__A	)



					def 	_lowercase					( self					:  Union[str, Any]							,		__A					:  str							,		__A					:  Optional[str] = None	):
						snake_case__			:				str       =  self._tokenizer.model.save(__A							,		name=__A	)
						return tuple(__A	)



					def 	_lowercase					( self					:  Union[str, Any]							,		__A					:  int							,		__A					:  str=None	):
						snake_case__			:				int       =  [self.bos_token_id] + token_ids_a + [self.eos_token_id]
						if token_ids_a is None:
							return output

						return output + [self.eos_token_id] + token_ids_a + [self.eos_token_id]



					def 	_lowercase					( self					:  Any							,		__A					:  List[int]							,		__A					:  Optional[List[int]] = None	):
						snake_case__			:				Optional[Any]       =  [self.sep_token_id]
						snake_case__			:				Any       =  [self.cls_token_id]

						if token_ids_a is None:
							return len(cls + token_ids_a + sep	) * [0]
						return len(cls + token_ids_a + sep + sep + token_ids_a + sep	) * [0]






					def 	_lowercase					( self					:  str							,		__A					:  Union[Dict[str, EncodedInput], BatchEncoding]							,		__A					:  Optional[int] = None							,		__A					:  PaddingStrategy = PaddingStrategy.DO_NOT_PAD							,		__A					:  Optional[int] = None							,		__A					:  Optional[bool] = None							,		):
						snake_case__			:				Dict       =  super()._pad(
						    encoded_inputs=__A							,		max_length=__A							,		padding_strategy=__A							,		pad_to_multiple_of=__A							,		return_attention_mask=__A							,		)

						# Load from model defaults
						if return_attention_mask is None:
							snake_case__			:				Optional[Any]       =  "attention_mask" in self.model_input_names

						if return_attention_mask and "global_attention_mask" in encoded_inputs:
							snake_case__			:				List[Any]       =  encoded_inputs[self.model_input_names[0]]
							# `global_attention_mask` need to have the same length as other (sequential) inputs.
							snake_case__			:				Union[str, Any]       =  len(encoded_inputs["global_attention_mask"]	) != len(__A	)

							if needs_to_be_padded:
								snake_case__			:				Dict       =  len(__A	) - len(encoded_inputs["global_attention_mask"]	)

								if self.padding_side == "right":
									# Use `-1` since `0` in `global_attention_mask` means `local attention` instead of `not to attend`
									snake_case__			:				Optional[Any]       =  (
									    encoded_inputs["global_attention_mask"] + [-1] * difference
									)
								elif self.padding_side == "left":
									snake_case__			:				Optional[Any]       =  [-1] * difference + encoded_inputs[
									    "global_attention_mask"
									]
								else:
									raise ValueError("Invalid padding strategy:" + str(self.padding_side	)	)

						return encoded_inputs



