






"""simple docstring"""





class 	a		:






    def __init__(			self	:				int   )	->						List[str]:
      lowerCamelCase_       	=		0
      lowerCamelCase_       	=		0
      lowerCamelCase_       	=		{}






    def   UpperCamelCase			(			self	:				str			,		__SCREAMING_SNAKE_CASE	:				Any   )	->						Tuple:
      if vertex not in self.adjacency:
        lowerCamelCase_       	=		{}
        self.num_vertices += 1






    def   UpperCamelCase			(			self	:				Union[str, Any]			,		__SCREAMING_SNAKE_CASE	:				str			,		__SCREAMING_SNAKE_CASE	:				Any			,		__SCREAMING_SNAKE_CASE	:				str   )	->						Union[str, Any]:
      self.add_vertex(__SCREAMING_SNAKE_CASE   )
      self.add_vertex(__SCREAMING_SNAKE_CASE   )

      if head == tail:
        return

      lowerCamelCase_       	=		weight
      lowerCamelCase_       	=		weight






    def   UpperCamelCase			(			self	:				Optional[int]   )	->						Tuple:
      lowerCamelCase_       	=		self.get_edges()
      for edge in edges:
        lowerCamelCase_	,      lowerCamelCase_	,      lowerCamelCase_       	=		edge
        edges.remove((tail, head, weight)   )
      for i in range(len(__SCREAMING_SNAKE_CASE   )   ):
        lowerCamelCase_       	=		list(edges[i]   )

      edges.sort(key=lambda __SCREAMING_SNAKE_CASE   : e[2]   )
      for i in range(len(__SCREAMING_SNAKE_CASE   ) - 1   ):
        if edges[i][2] >= edges[i + 1][2]:
          lowerCamelCase_       	=		edges[i][2] + 1
      for edge in edges:
        lowerCamelCase_	,      lowerCamelCase_	,      lowerCamelCase_       	=		edge
        lowerCamelCase_       	=		weight
        lowerCamelCase_       	=		weight






    def __str__(			self	:				Any   )	->						Union[str, Any]:
      lowerCamelCase_       	=		''
      for tail in self.adjacency:
        for head in self.adjacency[tail]:
          lowerCamelCase_       	=		self.adjacency[head][tail]
          string += F'''{head} -> {tail} == {weight}\n'''
      return string.rstrip('\n'   )






    def   UpperCamelCase			(			self	:				Any   )	->						Optional[int]:
      lowerCamelCase_       	=		[]
      for tail in self.adjacency:
        for head in self.adjacency[tail]:
          output.append((tail, head, self.adjacency[head][tail])   )
      return output






    def   UpperCamelCase			(			self	:				Optional[int]   )	->						Optional[int]:
      return self.adjacency.keys()






    @staticmethod
    def   UpperCamelCase			(			__SCREAMING_SNAKE_CASE	:				List[Any]=None			,		__SCREAMING_SNAKE_CASE	:				Optional[Any]=None   )	->						Dict:
      lowerCamelCase_       	=		Graph()
      if vertices is None:
        lowerCamelCase_       	=		[]
      if edges is None:
        lowerCamelCase_       	=		[]
      for vertex in vertices:
        g.add_vertex(__SCREAMING_SNAKE_CASE   )
      for edge in edges:
        g.add_edge(*__SCREAMING_SNAKE_CASE   )
      return g






    class 	a		:






        def __init__(			self	:				List[str]   )	->						Dict:
          lowerCamelCase_       	=		{}
          lowerCamelCase_       	=		{}






        def __len__(			self	:				Dict   )	->						Any:
          return len(self.parent   )






        def   UpperCamelCase			(			self	:				Tuple			,		__SCREAMING_SNAKE_CASE	:				Dict   )	->						Tuple:
          if item in self.parent:
            return self.find(__SCREAMING_SNAKE_CASE   )

          lowerCamelCase_       	=		item
          lowerCamelCase_       	=		0
          return item






        def   UpperCamelCase			(			self	:				Tuple			,		__SCREAMING_SNAKE_CASE	:				Optional[Any]   )	->						str:
          if item not in self.parent:
            return self.make_set(__SCREAMING_SNAKE_CASE   )
          if item != self.parent[item]:
            lowerCamelCase_       	=		self.find(self.parent[item]   )
          return self.parent[item]






        def   UpperCamelCase			(			self	:				Optional[int]			,		__SCREAMING_SNAKE_CASE	:				int			,		__SCREAMING_SNAKE_CASE	:				int   )	->						Any:
          lowerCamelCase_       	=		self.find(__SCREAMING_SNAKE_CASE   )
          lowerCamelCase_       	=		self.find(__SCREAMING_SNAKE_CASE   )

          if roota == roota:
            return roota

          if self.rank[roota] > self.rank[roota]:
            lowerCamelCase_       	=		roota
            return roota

          if self.rank[roota] < self.rank[roota]:
            lowerCamelCase_       	=		roota
            return roota

          if self.rank[roota] == self.rank[roota]:
            self.rank[roota] += 1
            lowerCamelCase_       	=		roota
            return roota
          return None






    @staticmethod
    def   UpperCamelCase			(			__SCREAMING_SNAKE_CASE	:				List[str]   )	->						Union[str, Any]:
      lowerCamelCase_       	=		graph.num_vertices

      lowerCamelCase_       	=		Graph.UnionFind()
      lowerCamelCase_       	=		[]
      while num_components > 1:
        lowerCamelCase_       	=		{}
        for vertex in graph.get_vertices():
          lowerCamelCase_       	=		-1

        lowerCamelCase_       	=		graph.get_edges()
        for edge in edges:
          lowerCamelCase_	,      lowerCamelCase_	,      lowerCamelCase_       	=		edge
          edges.remove((tail, head, weight)   )
        for edge in edges:
          lowerCamelCase_	,      lowerCamelCase_	,      lowerCamelCase_       	=		edge
          lowerCamelCase_       	=		union_find.find(__SCREAMING_SNAKE_CASE   )
          lowerCamelCase_       	=		union_find.find(__SCREAMING_SNAKE_CASE   )
          if seta != seta:
            if cheap_edge[seta] == -1 or cheap_edge[seta][2] > weight:
              lowerCamelCase_       	=		[head, tail, weight]

            if cheap_edge[seta] == -1 or cheap_edge[seta][2] > weight:
              lowerCamelCase_       	=		[head, tail, weight]
        for vertex in cheap_edge:
          if cheap_edge[vertex] != -1:
            lowerCamelCase_	,      lowerCamelCase_	,      lowerCamelCase_       	=		cheap_edge[vertex]
            if union_find.find(__SCREAMING_SNAKE_CASE   ) != union_find.find(__SCREAMING_SNAKE_CASE   ):
              union_find.union(__SCREAMING_SNAKE_CASE			,		__SCREAMING_SNAKE_CASE   )
              mst_edges.append(cheap_edge[vertex]   )
              lowerCamelCase_       	=		num_components - 1
      lowerCamelCase_       	=		Graph.build(edges=__SCREAMING_SNAKE_CASE   )
      return mst












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


_SCREAMING_SNAKE_CASE   :							int          =		re.compile(R'''\s+''')







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







def 	lowerCamelCase__			(				_lowerCamelCase			:						Tuple					) ->							int:
  lowerCamelCase_       	=		[len(_lowerCamelCase					) for line in example['content'].splitlines()]
  return {"line_mean": np.mean(_lowerCamelCase					), "line_max": max(_lowerCamelCase					)}







def 	lowerCamelCase__			(				_lowerCamelCase			:						List[str]					) ->							int:
  lowerCamelCase_       	=		np.mean([c.isalnum() for c in example['content']]					)
  return {"alpha_frac": alpha_frac}







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







def 	lowerCamelCase__			(				_lowerCamelCase			:						Any						,							_lowerCamelCase			:						Optional[Any]=5					) ->							int:
  lowerCamelCase_       	=		['auto-generated', 'autogenerated', 'automatically generated']
  lowerCamelCase_       	=		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 	lowerCamelCase__			(				_lowerCamelCase			:						Tuple						,							_lowerCamelCase			:						Dict=5						,							_lowerCamelCase			:						List[str]=0.05					) ->							Tuple:
  lowerCamelCase_       	=		['unit tests', 'test file', 'configuration file']
  lowerCamelCase_       	=		example['content'].splitlines()
  lowerCamelCase_       	=		0
  lowerCamelCase_       	=		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_       	=		example['content'].count('\n'					)
  lowerCamelCase_       	=		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__			(				_lowerCamelCase			:						Any					) ->							List[str]:
  lowerCamelCase_       	=		['def ', 'class ', 'for ', 'while ']
  lowerCamelCase_       	=		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__			(				_lowerCamelCase			:						Dict						,							_lowerCamelCase			:						Dict=4					) ->							Optional[Any]:
  lowerCamelCase_       	=		example['content'].splitlines()
  lowerCamelCase_       	=		0
  for line in lines:
    counter += line.lower().count('='					)
    if counter > minimum:
      return {"has_few_assignments": False}
  return {"has_few_assignments": True}







def 	lowerCamelCase__			(				_lowerCamelCase			:						Dict					) ->							List[str]:
  lowerCamelCase_       	=		tokenizer(example['content']						,							truncation=_lowerCamelCase					)['input_ids']
  lowerCamelCase_       	=		len(example['content']					) / len(_lowerCamelCase					)
  return {"ratio": ratio}







def 	lowerCamelCase__			(				_lowerCamelCase			:						Optional[Any]					) ->							List[Any]:
  lowerCamelCase_       	=		{}
  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 	lowerCamelCase__			(				_lowerCamelCase			:						Any						,							_lowerCamelCase			:						Dict						,							_lowerCamelCase			:						Optional[Any]					) ->							Any:
  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 	lowerCamelCase__			(				_lowerCamelCase			:						str					) ->							int:
  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
_SCREAMING_SNAKE_CASE   :							Optional[Any]          =		HfArgumentParser(PreprocessingArguments)
_SCREAMING_SNAKE_CASE   :							Any          =		parser.parse_args()
if args.num_workers is None:
  _SCREAMING_SNAKE_CASE   :							List[str]          =		multiprocessing.cpu_count()
_SCREAMING_SNAKE_CASE   :							str          =		AutoTokenizer.from_pretrained(args.tokenizer_dir)

# Load dataset
_SCREAMING_SNAKE_CASE   :							Optional[Any]          =		time.time()
_SCREAMING_SNAKE_CASE   :							int          =		load_dataset(args.dataset_name, split='''train''')
print(F'''Time to load dataset: {time.time()-t_start:.2f}''')

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

# Deduplicate hashes
_SCREAMING_SNAKE_CASE   :							List[Any]          =		set(ds.unique('''hash'''))
_SCREAMING_SNAKE_CASE   :							Optional[int]          =		len(uniques) / len(ds)
print(F'''Fraction of duplicates: {1-frac:.2%}''')

# Deduplicate data and apply heuristics
_SCREAMING_SNAKE_CASE   :							Dict          =		time.time()
_SCREAMING_SNAKE_CASE   :							Tuple          =		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:
  _SCREAMING_SNAKE_CASE   :							Optional[Any]          =		time.time()
  _SCREAMING_SNAKE_CASE       ,				_SCREAMING_SNAKE_CASE   :							Union[str, 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
_SCREAMING_SNAKE_CASE   :							Dict          =		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)

_SCREAMING_SNAKE_CASE   :							Optional[Any]          =		output_dir / '''data'''
data_dir.mkdir(exist_ok=True)

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




