
'''simple docstring'''






import copy
from typing import Any, Dict, List, Optional, Union

import numpy as np

from ...audio_utils import mel_filter_bank, spectrogram, window_function
from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
from ...feature_extraction_utils import BatchFeature
from ...utils import TensorType, logging


UpperCAmelCase_       :   str   =       logging.get_logger(__name__)

class  UpperCAmelCase__						(       _UpperCamelCase ):
    lowerCAmelCase_				    =		['input_features']



    def __init__(	self      :       Dict,__A      :       int=8_0,__A      :       Union[str, Any]=1_6_0_0_0,__A      :       Optional[int]=1_6_0,__A      :       Optional[Any]=3_0,__A      :       Dict=4_0_0,__A      :       Optional[Any]=0.0,__A      :       int=False,**__A      :       Tuple,):
          super().__init__(
              feature_size=__A,sampling_rate=__A,padding_value=__A,return_attention_mask=__A,**__A,)
          _lowerCamelCase				:   List[str]								=   n_fft
          _lowerCamelCase				:   Any								=   hop_length
          _lowerCamelCase				:   Dict								=   chunk_length
          _lowerCamelCase				:   Union[str, Any]								=   chunk_length * sampling_rate
          _lowerCamelCase				:   Any								=   self.n_samples // hop_length
          _lowerCamelCase				:   List[Any]								=   sampling_rate
          _lowerCamelCase				:   List[str]								=   mel_filter_bank(
              num_frequency_bins=1 + n_fft // 2,num_mel_filters=__A,min_frequency=0.0,max_frequency=8_0_0_0.0,sampling_rate=__A,norm="slaney",mel_scale="slaney",)



    def   lowerCamelCase_  (	self      :       Optional[Any],__A      :       np.array     ):
          _lowerCamelCase				:   Tuple								=   spectrogram(
              __A,window_function(self.n_fft,"hann"     ),frame_length=self.n_fft,hop_length=self.hop_length,power=2.0,mel_filters=self.mel_filters,log_mel="log10",)
          _lowerCamelCase				:   Tuple								=   log_spec[:, :-1]
          _lowerCamelCase				:   Any								=   np.maximum(__A,log_spec.max() - 8.0     )
          _lowerCamelCase				:   Optional[Any]								=   (log_spec + 4.0) / 4.0
          return log_spec



    @staticmethod
    # Copied from transformers.models.wav2vec2.feature_extraction_wav2vec2.Wav2Vec2FeatureExtractor.zero_mean_unit_var_norm
    def   lowerCamelCase_  (	__A      :       List[np.ndarray],__A      :       List[np.ndarray],__A      :       float = 0.0     ):
          if attention_mask is not None:
                _lowerCamelCase				:   List[str]								=   np.array(__A,np.intaa     )
                _lowerCamelCase				:   List[str]								=   []

                for vector, length in zip(__A,attention_mask.sum(-1     )     ):
                      _lowerCamelCase				:   List[Any]								=   (vector - vector[:length].mean()) / np.sqrt(vector[:length].var() + 1e-7     )
                      if length < normed_slice.shape[0]:
                            _lowerCamelCase				:   List[Any]								=   padding_value

                      normed_input_values.append(__A     )
          else:
                _lowerCamelCase				:   Any								=   [(x - x.mean()) / np.sqrt(x.var() + 1e-7     ) for x in input_values]

          return normed_input_values



    def __call__(	self      :       str,__A      :       Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]],__A      :       bool = True,__A      :       Optional[int] = None,__A      :       Optional[Union[str, TensorType]] = None,__A      :       Optional[bool] = None,__A      :       Optional[str] = "max_length",__A      :       Optional[int] = None,__A      :       Optional[int] = None,__A      :       Optional[bool] = None,**__A      :       List[Any],):

          if sampling_rate is not None:
                if sampling_rate != self.sampling_rate:
                      raise ValueError(
                          f'The model corresponding to this feature extractor: {self.__class__.__name__} was trained using a'
                          f' sampling rate of {self.sampling_rate}. Please make sure that the provided `raw_speech` input'
                          f' was sampled with {self.sampling_rate} and not {sampling_rate}.'     )
          else:
                logger.warning(
                    "It is strongly recommended to pass the `sampling_rate` argument to this function. "
                    "Failing to do so can result in silent errors that might be hard to debug."     )

          _lowerCamelCase				:   str								=   isinstance(__A,np.ndarray     ) and len(raw_speech.shape     ) > 1
          if is_batched_numpy and len(raw_speech.shape     ) > 2:
                raise ValueError(f'Only mono-channel audio is supported for input to {self}'     )
          _lowerCamelCase				:   Any								=   is_batched_numpy or (
              isinstance(__A,(list, tuple)     ) and (isinstance(raw_speech[0],(np.ndarray, tuple, list)     ))
          )

          if is_batched:
                _lowerCamelCase				:   Union[str, Any]								=   [np.asarray([speech],dtype=np.floataa     ).T for speech in raw_speech]
          elif not is_batched and not isinstance(__A,np.ndarray     ):
                _lowerCamelCase				:   Tuple								=   np.asarray(__A,dtype=np.floataa     )
          elif isinstance(__A,np.ndarray     ) and raw_speech.dtype is np.dtype(np.floataa     ):
                _lowerCamelCase				:   Optional[int]								=   raw_speech.astype(np.floataa     )

          # always return batch
          if not is_batched:
                _lowerCamelCase				:   Optional[Any]								=   [np.asarray([raw_speech]     ).T]

          _lowerCamelCase				:   List[str]								=   BatchFeature({"input_features": raw_speech}     )

          # convert into correct format for padding

          _lowerCamelCase				:   List[Any]								=   self.pad(
              __A,padding=__A,max_length=max_length if max_length else self.n_samples,truncation=__A,pad_to_multiple_of=__A,return_attention_mask=return_attention_mask or do_normalize,)

          # zero-mean and unit-variance normalization
          if do_normalize:
                _lowerCamelCase				:   Dict								=   self.zero_mean_unit_var_norm(
                    padded_inputs["input_features"],attention_mask=padded_inputs["attention_mask"],padding_value=self.padding_value,)
                _lowerCamelCase				:   Dict								=   np.stack(padded_inputs["input_features"],axis=0     )

          # make sure list is in array format
          _lowerCamelCase				:   Tuple								=   padded_inputs.get("input_features"     ).transpose(2,0,1     )

          _lowerCamelCase				:   str								=   [self._np_extract_fbank_features(__A     ) for waveform in input_features[0]]

          if isinstance(input_features[0],__A     ):
                _lowerCamelCase				:   Optional[int]								=   [np.asarray(__A,dtype=np.floataa     ) for feature in input_features]
          else:
                _lowerCamelCase				:   List[Any]								=   input_features

          if return_attention_mask:
                # rescale from sample (48000) to feature (3000)
                _lowerCamelCase				:   Optional[int]								=   padded_inputs["attention_mask"][:, :: self.hop_length]

          if return_tensors is not None:
                _lowerCamelCase				:   str								=   padded_inputs.convert_to_tensors(__A     )

          return padded_inputs



    def   lowerCamelCase_  (	self      :       Dict     ):
          _lowerCamelCase				:   Dict								=   copy.deepcopy(self.__dict__     )
          _lowerCamelCase				:   int								=   self.__class__.__name__
          if "mel_filters" in output:
                del output["mel_filters"]
          return output

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


UpperCAmelCase_       :   Any   =       re.compile(R'\s+')







def 		A_       (	_lowerCAmelCase   :    List[str]							):



      """simple docstring"""




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







def 		A_       (	_lowerCAmelCase   :    Union[str, Any]							):



      """simple docstring"""




      _lowerCamelCase				:   Union[str, Any]								=   [len(_lowerCAmelCase							) for line in example["content"].splitlines()]
      return {"line_mean": np.mean(_lowerCAmelCase							), "line_max": max(_lowerCAmelCase							)}







def 		A_       (	_lowerCAmelCase   :    int							):



      """simple docstring"""




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







def 		A_       (	_lowerCAmelCase   :    str				,	_lowerCAmelCase   :    Tuple							):



      """simple docstring"""




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







def 		A_       (	_lowerCAmelCase   :    Optional[int]				,	_lowerCAmelCase   :    Dict=5							):



      """simple docstring"""




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







def 		A_       (	_lowerCAmelCase   :    List[str]				,	_lowerCAmelCase   :    Optional[Any]=5				,	_lowerCAmelCase   :    List[Any]=0.0_5							):



      """simple docstring"""




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



      """simple docstring"""




      _lowerCamelCase				:   str								=   ["def ", "class ", "for ", "while "]
      _lowerCamelCase				:   List[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_       (	_lowerCAmelCase   :    Dict				,	_lowerCAmelCase   :    int=4							):



      """simple docstring"""




      _lowerCamelCase				:   List[Any]								=   example["content"].splitlines()
      _lowerCamelCase				:   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_       (	_lowerCAmelCase   :    List[str]							):



      """simple docstring"""




      _lowerCamelCase				:   List[Any]								=   tokenizer(example["content"]				,	truncation=_lowerCAmelCase							)["input_ids"]
      _lowerCamelCase				:   str								=   len(example["content"]							) / len(_lowerCAmelCase							)
      return {"ratio": ratio}







def 		A_       (	_lowerCAmelCase   :    Optional[Any]							):



      """simple docstring"""




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







def 		A_       (	_lowerCAmelCase   :    Union[str, Any]				,	_lowerCAmelCase   :    Optional[Any]				,	_lowerCAmelCase   :    str							):



      """simple docstring"""




      if not check_uniques(_lowerCAmelCase				,	_lowerCAmelCase							):
            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_       (	_lowerCAmelCase   :    int							):



      """simple docstring"""




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


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

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

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

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

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

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