
'''simple docstring'''







import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

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


__lowerCAmelCase         =    logging.get_logger(__name__)

__lowerCAmelCase         =    """▁"""

__lowerCAmelCase         =    {"""vocab_file""": """sentencepiece.bpe.model"""}

__lowerCAmelCase         =    {
    """vocab_file""": {
        """facebook/mbart-large-50-one-to-many-mmt""": (
            """https://huggingface.co/facebook/mbart-large-50-one-to-many-mmt/resolve/main/sentencepiece.bpe.model"""
        ),
    }
}

__lowerCAmelCase         =    {
    """facebook/mbart-large-50-one-to-many-mmt""": 1_0_2_4,
}

# fmt: off
__lowerCAmelCase         =    ["""ar_AR""", """cs_CZ""", """de_DE""", """en_XX""", """es_XX""", """et_EE""", """fi_FI""", """fr_XX""", """gu_IN""", """hi_IN""", """it_IT""", """ja_XX""", """kk_KZ""", """ko_KR""", """lt_LT""", """lv_LV""", """my_MM""", """ne_NP""", """nl_XX""", """ro_RO""", """ru_RU""", """si_LK""", """tr_TR""", """vi_VN""", """zh_CN""", """af_ZA""", """az_AZ""", """bn_IN""", """fa_IR""", """he_IL""", """hr_HR""", """id_ID""", """ka_GE""", """km_KH""", """mk_MK""", """ml_IN""", """mn_MN""", """mr_IN""", """pl_PL""", """ps_AF""", """pt_XX""", """sv_SE""", """sw_KE""", """ta_IN""", """te_IN""", """th_TH""", """tl_XX""", """uk_UA""", """ur_PK""", """xh_ZA""", """gl_ES""", """sl_SI"""]







class UpperCAmelCase__	( lowercase__	):





  """simple docstring"""

  __UpperCAmelCase				:				Tuple   =					VOCAB_FILES_NAMES
  __UpperCAmelCase				:				Tuple   =					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
  __UpperCAmelCase				:				str   =					PRETRAINED_VOCAB_FILES_MAP
  __UpperCAmelCase				:				int   =					['''input_ids''', '''attention_mask''']

  __UpperCAmelCase				:				List[int]   =					[]
  __UpperCAmelCase				:				List[int]   =					[]


  def __init__(		self      :     Any							,_a      :     int							,_a      :     Union[str, Any]=None							,_a      :     Optional[int]=None							,_a      :     int="</s>"							,_a      :     Optional[Any]="</s>"							,_a      :     Dict="<s>"							,_a      :     List[Any]="<unk>"							,_a      :     Tuple="<pad>"							,_a      :     List[Any]="<mask>"							,_a      :     Optional[Dict[str, Any]] = None							,**_a      :     List[str]							,):


     '''simple docstring'''






     _a				:	Optional[Any]			=   AddedToken(_a							,lstrip=_a							,rstrip=_a       ) if isinstance(_a							,_a       ) else mask_token

     _a				:	List[str]			=   {} if sp_model_kwargs is None else sp_model_kwargs

     _a				:	Tuple			=   kwargs.get('additional_special_tokens'							,[]       )
     kwargs["additional_special_tokens"] += [
         code for code in FAIRSEQ_LANGUAGE_CODES if code not in kwargs["additional_special_tokens"]
     ]

     super().__init__(
         src_lang=_a							,tgt_lang=_a							,eos_token=_a							,unk_token=_a							,sep_token=_a							,cls_token=_a							,pad_token=_a							,mask_token=_a							,sp_model_kwargs=self.sp_model_kwargs							,**_a							,)

     _a				:	Tuple			=   spm.SentencePieceProcessor(**self.sp_model_kwargs       )
     self.sp_model.Load(str(_a       )       )
     _a				:	List[Any]			=   vocab_file

     # Original fairseq vocab and spm vocab must be "aligned":
     # Vocab    |    0    |    1    |   2    |    3    |  4  |  5  |  6  |   7   |   8   |  9
     # -------- | ------- | ------- | ------ | ------- | --- | --- | --- | ----- | ----- | ----
     # fairseq  | '<s>'   | '<pad>' | '</s>' | '<unk>' | ',' | '.' | '▁' | 's'   | '▁de' | '-'
     # spm      | '<unk>' | '<s>'   | '</s>' | ','     | '.' | '▁' | 's' | '▁de' | '-'   | '▁a'

     # Mimic fairseq token-to-id alignment for the first 4 token
     _a				:	Optional[int]			=   {'<s>': 0, '<pad>': 1, '</s>': 2, '<unk>': 3}

     # The first "real" token "," has position 4 in the original fairseq vocab and position 3 in the spm vocab
     _a				:	Union[str, Any]			=   1

     _a				:	Optional[Any]			=   len(self.sp_model       )
     _a				:	Tuple			=   {
         code: self.sp_model_size + i + self.fairseq_offset for i, code in enumerate(_a       )
     }
     _a				:	Any			=   {v: k for k, v in self.lang_code_to_id.items()}
     _a				:	Dict			=   len(self.sp_model       ) + len(self.lang_code_to_id       ) + self.fairseq_offset

     self.fairseq_tokens_to_ids.update(self.lang_code_to_id       )
     _a				:	int			=   {v: k for k, v in self.fairseq_tokens_to_ids.items()}

     _a				:	Any			=   src_lang if src_lang is not None else 'en_XX'
     _a				:	List[Any]			=   self.lang_code_to_id[self._src_lang]
     _a				:	Union[str, Any]			=   tgt_lang
     self.set_src_lang_special_tokens(self._src_lang       )


  @property
  def   __lowercase		(		self      :     Union[str, Any]       ):


     '''simple docstring'''






     return len(self.sp_model       ) + len(self.lang_code_to_id       ) + self.fairseq_offset + 1  # Plus 1 for the mask token


  @property
  def   __lowercase		(		self      :     int       ):


     '''simple docstring'''






     return self._src_lang


  @src_lang.setter
  def   __lowercase		(		self      :     Any							,_a      :     str       ):


     '''simple docstring'''






     _a				:	Optional[Any]			=   new_src_lang
     self.set_src_lang_special_tokens(self._src_lang       )


  def __getstate__(		self      :     Dict       ):


     '''simple docstring'''






     _a				:	Union[str, Any]			=   self.__dict__.copy()
     _a				:	int			=   None
     return state


  def __setstate__(		self      :     int							,_a      :     Dict       ):


     '''simple docstring'''






     _a				:	int			=   d

     # for backward compatibility
     if not hasattr(self							,'sp_model_kwargs'       ):
        _a				:	Dict			=   {}

     _a				:	Tuple			=   spm.SentencePieceProcessor(**self.sp_model_kwargs       )
     self.sp_model.Load(self.vocab_file       )


  def   __lowercase		(		self      :     Any       ):


     '''simple docstring'''






     _a				:	Dict			=   {self.convert_ids_to_tokens(_a       ): i for i in range(self.vocab_size       )}
     vocab.update(self.added_tokens_encoder       )
     return vocab


  def   __lowercase		(		self      :     str							,_a      :     str       ):


     '''simple docstring'''






     return self.sp_model.encode(_a							,out_type=_a       )


  def   __lowercase		(		self      :     str							,_a      :     str       ):


     '''simple docstring'''






     if token in self.fairseq_tokens_to_ids:
        return self.fairseq_tokens_to_ids[token]
     _a				:	Tuple			=   self.sp_model.PieceToId(_a       )

     # Need to return unknown token if the SP model returned 0
     return spm_id + self.fairseq_offset if spm_id else self.unk_token_id


  def   __lowercase		(		self      :     Optional[int]							,_a      :     int       ):


     '''simple docstring'''






     if index in self.fairseq_ids_to_tokens:
        return self.fairseq_ids_to_tokens[index]
     return self.sp_model.IdToPiece(index - self.fairseq_offset       )


  def   __lowercase		(		self      :     Dict							,_a      :     str       ):


     '''simple docstring'''






     _a				:	Union[str, Any]			=   []
     _a				:	Dict			=   ''
     _a				:	Any			=   False
     for token in tokens:
        # make sure that special tokens are not decoded using sentencepiece model
        if token in self.all_special_tokens:
           if not prev_is_special:
              out_string += " "
           out_string += self.sp_model.decode(_a       ) + token
           _a				:	Any			=   True
           _a				:	int			=   []
        else:
           current_sub_tokens.append(_a       )
           _a				:	int			=   False
     out_string += self.sp_model.decode(_a       )
     return out_string.strip()


  def   __lowercase		(		self      :     Optional[int]							,_a      :     str							,_a      :     Optional[str] = None       ):


     '''simple docstring'''






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

     if os.path.abspath(self.vocab_file       ) != os.path.abspath(_a       ) and os.path.isfile(self.vocab_file       ):
        copyfile(self.vocab_file							,_a       )
     elif not os.path.isfile(self.vocab_file       ):
        with open(_a							,'wb'       ) as fi:
           _a				:	Optional[int]			=   self.sp_model.serialized_model_proto()
           fi.write(_a       )

     return (out_vocab_file,)


  def   __lowercase		(		self      :     List[str]							,_a      :     List[int]							,_a      :     Optional[List[int]] = None							,_a      :     bool = False       ):


     '''simple docstring'''






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

     _a				:	Optional[Any]			=   [1] * len(self.prefix_tokens       )
     _a				:	List[Any]			=   [1] * len(self.suffix_tokens       )
     if token_ids_a is None:
        return prefix_ones + ([0] * len(_a       )) + suffix_ones
     return prefix_ones + ([0] * len(_a       )) + ([0] * len(_a       )) + suffix_ones


  def   __lowercase		(		self      :     Dict							,_a      :     List[int]							,_a      :     Optional[List[int]] = None       ):


     '''simple docstring'''






     if token_ids_a is None:
        return self.prefix_tokens + token_ids_a + self.suffix_tokens
     # We don't expect to process pairs, but leave the pair logic for API consistency
     return self.prefix_tokens + token_ids_a + token_ids_a + self.suffix_tokens


  def   __lowercase		(		self      :     Union[str, Any]							,_a      :     Optional[Any]							,_a      :     str							,_a      :     Optional[str]							,_a      :     Optional[str]							,**_a      :     str       ):


     '''simple docstring'''






     if src_lang is None or tgt_lang is None:
        raise ValueError('Translation requires a `src_lang` and a `tgt_lang` for this model'       )
     _a				:	str			=   src_lang
     _a				:	Union[str, Any]			=   self(_a							,add_special_tokens=_a							,return_tensors=_a							,**_a       )
     _a				:	Any			=   self.convert_tokens_to_ids(_a       )
     _a				:	int			=   tgt_lang_id
     return inputs


  def   __lowercase		(		self      :     Optional[Any]							,_a      :     List[str]							,_a      :     str = "en_XX"							,_a      :     Optional[List[str]] = None							,_a      :     str = "ro_RO"							,**_a      :     str							,):


     '''simple docstring'''






     _a				:	Dict			=   src_lang
     _a				:	List[Any]			=   tgt_lang
     return super().prepare_seqaseq_batch(_a							,_a							,**_a       )


  def   __lowercase		(		self      :     Tuple       ):


     '''simple docstring'''






     return self.set_src_lang_special_tokens(self.src_lang       )


  def   __lowercase		(		self      :     Dict       ):


     '''simple docstring'''






     return self.set_tgt_lang_special_tokens(self.tgt_lang       )


  def   __lowercase		(		self      :     Optional[int]							,_a      :     str       ):


     '''simple docstring'''






     _a				:	str			=   self.lang_code_to_id[src_lang]
     _a				:	List[Any]			=   [self.cur_lang_code_id]
     _a				:	Optional[int]			=   [self.eos_token_id]


  def   __lowercase		(		self      :     Optional[int]							,_a      :     str       ):


     '''simple docstring'''






     _a				:	int			=   self.lang_code_to_id[tgt_lang]
     _a				:	Optional[Any]			=   [self.cur_lang_code_id]
     _a				:	str			=   [self.eos_token_id]





'''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 					UpperCAmelCase_		(__a      :    Any ):






   """simple docstring"""

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





def 					UpperCAmelCase_		(__a      :    str ):






   """simple docstring"""

   _a				:	List[str]			=   [len(__a ) for line in example['content'].splitlines()]
   return {"line_mean": np.mean(__a ), "line_max": max(__a )}





def 					UpperCAmelCase_		(__a      :    str ):






   """simple docstring"""

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





def 					UpperCAmelCase_		(__a      :    Optional[int] ,	__a      :    Any ):






   """simple docstring"""

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





def 					UpperCAmelCase_		(__a      :    int ,	__a      :    Union[str, Any]=5 ):






   """simple docstring"""

   _a				:	Optional[int]			=   ['auto-generated', 'autogenerated', 'automatically generated']
   _a				:	List[str]			=   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 					UpperCAmelCase_		(__a      :    List[str] ,	__a      :    Dict=5 ,	__a      :    Tuple=0.05 ):






   """simple docstring"""

   _a				:	Optional[int]			=   ['unit tests', 'test file', 'configuration file']
   _a				:	int			=   example['content'].splitlines()
   _a				:	int			=   0
   _a				:	Dict			=   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
   _a				:	int			=   example['content'].count('\n' )
   _a				:	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 					UpperCAmelCase_		(__a      :    Optional[int] ):






   """simple docstring"""

   _a				:	List[str]			=   ['def ', 'class ', 'for ', 'while ']
   _a				:	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 					UpperCAmelCase_		(__a      :    int ,	__a      :    Any=4 ):






   """simple docstring"""

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





def 					UpperCAmelCase_		(__a      :    str ):






   """simple docstring"""

   _a				:	Optional[Any]			=   tokenizer(example['content'] ,	truncation=__a )['input_ids']
   _a				:	Optional[int]			=   len(example['content'] ) / len(__a )
   return {"ratio": ratio}





def 					UpperCAmelCase_		(__a      :    str ):






   """simple docstring"""

   _a				:	Dict			=   {}
   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 					UpperCAmelCase_		(__a      :    Any ,	__a      :    Any ,	__a      :    str ):






   """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 					UpperCAmelCase_		(__a      :    Union[str, Any] ):






   """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
__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}''')



