

import json
import os
from functools import lru_cache
from typing import List, Optional, Tuple

import regex as re

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


UpperCamelCase__		 =     logging.get_logger(__name__)


UpperCamelCase__		 =     {'''vocab_file''': '''vocab.json''', '''merges_file''': '''merges.txt'''}

UpperCamelCase__		 =     {
    '''vocab_file''': {
        '''allenai/longformer-base-4096''': '''https://huggingface.co/allenai/longformer-base-4096/resolve/main/vocab.json''',
        '''allenai/longformer-large-4096''': (
            '''https://huggingface.co/allenai/longformer-large-4096/resolve/main/vocab.json'''
        ),
        '''allenai/longformer-large-4096-finetuned-triviaqa''': (
            '''https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa/resolve/main/vocab.json'''
        ),
        '''allenai/longformer-base-4096-extra.pos.embd.only''': (
            '''https://huggingface.co/allenai/longformer-base-4096-extra.pos.embd.only/resolve/main/vocab.json'''
        ),
        '''allenai/longformer-large-4096-extra.pos.embd.only''': (
            '''https://huggingface.co/allenai/longformer-large-4096-extra.pos.embd.only/resolve/main/vocab.json'''
        ),
    },
    '''merges_file''': {
        '''allenai/longformer-base-4096''': '''https://huggingface.co/allenai/longformer-base-4096/resolve/main/merges.txt''',
        '''allenai/longformer-large-4096''': (
            '''https://huggingface.co/allenai/longformer-large-4096/resolve/main/merges.txt'''
        ),
        '''allenai/longformer-large-4096-finetuned-triviaqa''': (
            '''https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa/resolve/main/merges.txt'''
        ),
        '''allenai/longformer-base-4096-extra.pos.embd.only''': (
            '''https://huggingface.co/allenai/longformer-base-4096-extra.pos.embd.only/resolve/main/merges.txt'''
        ),
        '''allenai/longformer-large-4096-extra.pos.embd.only''': (
            '''https://huggingface.co/allenai/longformer-large-4096-extra.pos.embd.only/resolve/main/merges.txt'''
        ),
    },
}

UpperCamelCase__		 =     {
    '''allenai/longformer-base-4096''': 4096,
    '''allenai/longformer-large-4096''': 4096,
    '''allenai/longformer-large-4096-finetuned-triviaqa''': 4096,
    '''allenai/longformer-base-4096-extra.pos.embd.only''': 4096,
    '''allenai/longformer-large-4096-extra.pos.embd.only''': 4096,
}







@lru_cache()
# Copied from transformers.models.roberta.tokenization_roberta.bytes_to_unicode
def 						lowerCamelCase__      ( ):


       """simple docstring"""




       __snake_case       	=							(
           list(range(ord("""!"""		)	,ord("""~"""		) + 1		)		) + list(range(ord("""¡"""		)	,ord("""¬"""		) + 1		)		) + list(range(ord("""®"""		)	,ord("""ÿ"""		) + 1		)		)
       )
       __snake_case       	=							bs[:]
       __snake_case       	=							0
       for b in range(2**8		):
              if b not in bs:
                     bs.append(__A		)
                     cs.append(2**8 + n		)
                     n += 1
       __snake_case       	=							[chr(__A		) for n in cs]
       return dict(zip(__A	,__A		)		)







def 						lowerCamelCase__      ( __A    :List[Any]		):


       """simple docstring"""




       __snake_case       	=							set()
       __snake_case       	=							word[0]
       for char in word[1:]:
              pairs.add((prev_char, char)		)
              __snake_case       	=							char
       return pairs




class     __snake_case    (				snake_case__  ):
       """simple docstring"""







       __SCREAMING_SNAKE_CASE				    =							VOCAB_FILES_NAMES
       __SCREAMING_SNAKE_CASE				    =							PRETRAINED_VOCAB_FILES_MAP
       __SCREAMING_SNAKE_CASE				    =							PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
       __SCREAMING_SNAKE_CASE				    =							["input_ids", "attention_mask"]


       def __init__(		self      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase="replace"      , _UpperCamelCase="<s>"      , _UpperCamelCase="</s>"      , _UpperCamelCase="</s>"      , _UpperCamelCase="<s>"      , _UpperCamelCase="<unk>"      , _UpperCamelCase="<pad>"      , _UpperCamelCase="<mask>"      , _UpperCamelCase=False      , **_UpperCamelCase      , )      ->							List[str]:






              """simple docstring"""


              __snake_case       	=							AddedToken(_UpperCamelCase      , lstrip=_UpperCamelCase      , rstrip=_UpperCamelCase   ) if isinstance(_UpperCamelCase      , _UpperCamelCase   ) else bos_token
              __snake_case       	=							AddedToken(_UpperCamelCase      , lstrip=_UpperCamelCase      , rstrip=_UpperCamelCase   ) if isinstance(_UpperCamelCase      , _UpperCamelCase   ) else eos_token
              __snake_case       	=							AddedToken(_UpperCamelCase      , lstrip=_UpperCamelCase      , rstrip=_UpperCamelCase   ) if isinstance(_UpperCamelCase      , _UpperCamelCase   ) else sep_token
              __snake_case       	=							AddedToken(_UpperCamelCase      , lstrip=_UpperCamelCase      , rstrip=_UpperCamelCase   ) if isinstance(_UpperCamelCase      , _UpperCamelCase   ) else cls_token
              __snake_case       	=							AddedToken(_UpperCamelCase      , lstrip=_UpperCamelCase      , rstrip=_UpperCamelCase   ) if isinstance(_UpperCamelCase      , _UpperCamelCase   ) else unk_token
              __snake_case       	=							AddedToken(_UpperCamelCase      , lstrip=_UpperCamelCase      , rstrip=_UpperCamelCase   ) if isinstance(_UpperCamelCase      , _UpperCamelCase   ) else pad_token

              # Mask token behave like a normal word, i.e. include the space before it
              __snake_case       	=							AddedToken(_UpperCamelCase      , lstrip=_UpperCamelCase      , rstrip=_UpperCamelCase   ) if isinstance(_UpperCamelCase      , _UpperCamelCase   ) else mask_token

              super().__init__(
                  errors=_UpperCamelCase      , bos_token=_UpperCamelCase      , eos_token=_UpperCamelCase      , unk_token=_UpperCamelCase      , sep_token=_UpperCamelCase      , cls_token=_UpperCamelCase      , pad_token=_UpperCamelCase      , mask_token=_UpperCamelCase      , add_prefix_space=_UpperCamelCase      , **_UpperCamelCase      , )

              with open(_UpperCamelCase      , encoding="""utf-8"""   ) as vocab_handle:
                     __snake_case       	=							json.load(_UpperCamelCase   )
              __snake_case       	=							{v: k for k, v in self.encoder.items()}
              __snake_case       	=							errors  # how to handle errors in decoding
              __snake_case       	=							bytes_to_unicode()
              __snake_case       	=							{v: k for k, v in self.byte_encoder.items()}
              with open(_UpperCamelCase      , encoding="""utf-8"""   ) as merges_handle:
                     __snake_case       	=							merges_handle.read().split("""\n"""   )[1:-1]
              __snake_case       	=							[tuple(merge.split()   ) for merge in bpe_merges]
              __snake_case       	=							dict(zip(_UpperCamelCase      , range(len(_UpperCamelCase   )   )   )   )
              __snake_case       	=							{}
              __snake_case       	=							add_prefix_space

              # Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
              __snake_case       	=							re.compile(r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""   )


       @property
       def 			a		(		self   )      ->							Union[str, Any]:






              """simple docstring"""


              return len(self.encoder   )


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






              """simple docstring"""


              return dict(self.encoder      , **self.added_tokens_encoder   )


       def 			a		(		self      , _UpperCamelCase   )      ->							Union[str, Any]:






              """simple docstring"""


              if token in self.cache:
                     return self.cache[token]
              __snake_case       	=							tuple(_UpperCamelCase   )
              __snake_case       	=							get_pairs(_UpperCamelCase   )

              if not pairs:
                     return token

              while True:
                     __snake_case       	=							min(_UpperCamelCase      , key=lambda _UpperCamelCase   : self.bpe_ranks.get(_UpperCamelCase      , float("""inf"""   )   )   )
                     if bigram not in self.bpe_ranks:
                            break
                     __snake_case	,    __snake_case       	=							bigram
                     __snake_case       	=							[]
                     __snake_case       	=							0
                     while i < len(_UpperCamelCase   ):
                            try:
                                   __snake_case       	=							word.index(_UpperCamelCase      , _UpperCamelCase   )
                            except ValueError:
                                   new_word.extend(word[i:]   )
                                   break
                            else:
                                   new_word.extend(word[i:j]   )
                                   __snake_case       	=							j

                            if word[i] == first and i < len(_UpperCamelCase   ) - 1 and word[i + 1] == second:
                                   new_word.append(first + second   )
                                   i += 2
                            else:
                                   new_word.append(word[i]   )
                                   i += 1
                     __snake_case       	=							tuple(_UpperCamelCase   )
                     __snake_case       	=							new_word
                     if len(_UpperCamelCase   ) == 1:
                            break
                     else:
                            __snake_case       	=							get_pairs(_UpperCamelCase   )
              __snake_case       	=							""" """.join(_UpperCamelCase   )
              __snake_case       	=							word
              return word


       def 			a		(		self      , _UpperCamelCase   )      ->							Optional[Any]:






              """simple docstring"""


              __snake_case       	=							[]
              for token in re.findall(self.pat      , _UpperCamelCase   ):
                     __snake_case       	=							"""""".join(
                         self.byte_encoder[b] for b in token.encode("""utf-8"""   )   )  # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
                     bpe_tokens.extend(bpe_token for bpe_token in self.bpe(_UpperCamelCase   ).split(""" """   )   )
              return bpe_tokens


       def 			a		(		self      , _UpperCamelCase   )      ->							Optional[int]:






              """simple docstring"""


              return self.encoder.get(_UpperCamelCase      , self.encoder.get(self.unk_token   )   )


       def 			a		(		self      , _UpperCamelCase   )      ->							List[Any]:






              """simple docstring"""


              return self.decoder.get(_UpperCamelCase   )


       def 			a		(		self      , _UpperCamelCase   )      ->							Tuple:






              """simple docstring"""


              __snake_case       	=							"""""".join(_UpperCamelCase   )
              __snake_case       	=							bytearray([self.byte_decoder[c] for c in text]   ).decode("""utf-8"""      , errors=self.errors   )
              return text


       def 			a		(		self      , _UpperCamelCase      , _UpperCamelCase = None   )      ->							Tuple[str]:






              """simple docstring"""


              if not os.path.isdir(_UpperCamelCase   ):
                     logger.error(F'Vocabulary path ({save_directory}) should be a directory'   )
                     return
              __snake_case       	=							os.path.join(
                  _UpperCamelCase      , (filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]   )
              __snake_case       	=							os.path.join(
                  _UpperCamelCase      , (filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""merges_file"""]   )

              with open(_UpperCamelCase      , """w"""      , encoding="""utf-8"""   ) as f:
                     f.write(json.dumps(self.encoder      , indent=2      , sort_keys=_UpperCamelCase      , ensure_ascii=_UpperCamelCase   ) + """\n"""   )

              __snake_case       	=							0
              with open(_UpperCamelCase      , """w"""      , encoding="""utf-8"""   ) as writer:
                     writer.write("""#version: 0.2\n"""   )
                     for bpe_tokens, token_index in sorted(self.bpe_ranks.items()      , key=lambda _UpperCamelCase   : kv[1]   ):
                            if index != token_index:
                                   logger.warning(
                                       F'Saving vocabulary to {merge_file}: BPE merge indices are not consecutive.'
                                       """ Please check that the tokenizer is not corrupted!"""   )
                                   __snake_case       	=							token_index
                            writer.write(""" """.join(_UpperCamelCase   ) + """\n"""   )
                            index += 1

              return vocab_file, merge_file


       def 			a		(		self      , _UpperCamelCase      , _UpperCamelCase = None   )      ->							List[int]:






              """simple docstring"""


              if token_ids_a is None:
                     return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
              __snake_case       	=							[self.cls_token_id]
              __snake_case       	=							[self.sep_token_id]
              return cls + token_ids_a + sep + sep + token_ids_a + sep


       def 			a		(		self      , _UpperCamelCase      , _UpperCamelCase = None      , _UpperCamelCase = False   )      ->							List[int]:






              """simple docstring"""


              if already_has_special_tokens:
                     return super().get_special_tokens_mask(
                         token_ids_a=_UpperCamelCase      , token_ids_a=_UpperCamelCase      , already_has_special_tokens=_UpperCamelCase   )

              if token_ids_a is None:
                     return [1] + ([0] * len(_UpperCamelCase   )) + [1]
              return [1] + ([0] * len(_UpperCamelCase   )) + [1, 1] + ([0] * len(_UpperCamelCase   )) + [1]


       def 			a		(		self      , _UpperCamelCase      , _UpperCamelCase = None   )      ->							List[int]:






              """simple docstring"""


              __snake_case       	=							[self.sep_token_id]
              __snake_case       	=							[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 			a		(		self      , _UpperCamelCase      , _UpperCamelCase=False      , **_UpperCamelCase   )      ->							Union[str, Any]:






              """simple docstring"""


              __snake_case       	=							kwargs.pop("""add_prefix_space"""      , self.add_prefix_space   )
              if (is_split_into_words or add_prefix_space) and (len(_UpperCamelCase   ) > 0 and not text[0].isspace()):
                     __snake_case       	=							""" """ + text
              return (text, kwargs)




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


UpperCamelCase__		 =     re.compile(r'''\s+''')







def 						lowerCamelCase__      ( __A    :int		):


       """simple docstring"""




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







def 						lowerCamelCase__      ( __A    :Any		):


       """simple docstring"""




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







def 						lowerCamelCase__      ( __A    :Any		):


       """simple docstring"""




       __snake_case       	=							np.mean([c.isalnum() for c in example["""content"""]]		)
       return {"alpha_frac": alpha_frac}







def 						lowerCamelCase__      ( __A    :str	,__A    :int		):


       """simple docstring"""




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







def 						lowerCamelCase__      ( __A    :Optional[int]	,__A    :int=5		):


       """simple docstring"""




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







def 						lowerCamelCase__      ( __A    :int	,__A    :int=5	,__A    :str=0.05		):


       """simple docstring"""




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


       """simple docstring"""




       __snake_case       	=							["""def """, """class """, """for """, """while """]
       __snake_case       	=							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__      ( __A    :Dict	,__A    :List[str]=4		):


       """simple docstring"""




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







def 						lowerCamelCase__      ( __A    :str		):


       """simple docstring"""




       __snake_case       	=							tokenizer(example["""content"""]	,truncation=__A		)["""input_ids"""]
       __snake_case       	=							len(example["""content"""]		) / len(__A		)
       return {"ratio": ratio}







def 						lowerCamelCase__      ( __A    :Optional[Any]		):


       """simple docstring"""




       __snake_case       	=							{}
       results.update(get_hash(__A		)		)
       results.update(line_stats(__A		)		)
       results.update(alpha_stats(__A		)		)
       results.update(char_token_ratio(__A		)		)
       results.update(is_autogenerated(__A		)		)
       results.update(is_config_or_test(__A		)		)
       results.update(has_no_keywords(__A		)		)
       results.update(has_few_assignments(__A		)		)
       return results







def 						lowerCamelCase__      ( __A    :int	,__A    :str	,__A    :Dict		):


       """simple docstring"""




       if not check_uniques(__A	,__A		):
              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__      ( __A    :str		):


       """simple docstring"""




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


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

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

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

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

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

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

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

