"""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__  		=						re.compile(r"""\s+""")






def     __UpperCAmelCase					(			lowercase		):







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






def     __UpperCAmelCase					(			lowercase		):







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






def     __UpperCAmelCase					(			lowercase		):







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






def     __UpperCAmelCase					(			lowercase   ,lowercase		):







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






def     __UpperCAmelCase					(			lowercase   ,lowercase=5		):







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






def     __UpperCAmelCase					(			lowercase   ,lowercase=5   ,lowercase=0.05		):







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







		"""simple docstring"""
		_UpperCAmelCase						 =      ["""def """, """class """, """for """, """while """]
		_UpperCAmelCase						 =      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					(			lowercase   ,lowercase=4		):







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






def     __UpperCAmelCase					(			lowercase		):







		"""simple docstring"""
		_UpperCAmelCase						 =      tokenizer(example["""content"""]   ,truncation=lowercase		)["""input_ids"""]
		_UpperCAmelCase						 =      len(example["""content"""]		) / len(lowercase		)
		return {"ratio": ratio}






def     __UpperCAmelCase					(			lowercase		):







		"""simple docstring"""
		_UpperCAmelCase						 =      {}
		results.update(get_hash(lowercase		)		)
		results.update(line_stats(lowercase		)		)
		results.update(alpha_stats(lowercase		)		)
		results.update(char_token_ratio(lowercase		)		)
		results.update(is_autogenerated(lowercase		)		)
		results.update(is_config_or_test(lowercase		)		)
		results.update(has_no_keywords(lowercase		)		)
		results.update(has_few_assignments(lowercase		)		)
		return results






def     __UpperCAmelCase					(			lowercase   ,lowercase   ,lowercase		):







		"""simple docstring"""
		if not check_uniques(lowercase   ,lowercase		):
				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					(			lowercase		):







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


# 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}''')

"""simple docstring"""


import torch
from torch import nn


class        a				(		nn.Module   ):


    def __init__(					self	:  List[Any]   ,				__lowerCAmelCase	:  Optional[Any]   ,				__lowerCAmelCase	:  str   ,				__lowerCAmelCase	:  List[Any]   ,				__lowerCAmelCase	:  int   ,				__lowerCAmelCase	:  List[Any]=1   ,				__lowerCAmelCase	:  Tuple=False					):
      super().__init__()

      _UpperCAmelCase						 =      n_token
      _UpperCAmelCase						 =      d_embed
      _UpperCAmelCase						 =      d_proj

      _UpperCAmelCase						 =      cutoffs + [n_token]
      _UpperCAmelCase						 =      [0] + self.cutoffs
      _UpperCAmelCase						 =      div_val

      _UpperCAmelCase						 =      self.cutoffs[0]
      _UpperCAmelCase						 =      len(self.cutoffs					) - 1
      _UpperCAmelCase						 =      self.shortlist_size + self.n_clusters

      if self.n_clusters > 0:
        _UpperCAmelCase						 =      nn.Parameter(torch.zeros(self.n_clusters   ,				self.d_embed					)					)
        _UpperCAmelCase						 =      nn.Parameter(torch.zeros(self.n_clusters					)					)

      _UpperCAmelCase						 =      nn.ModuleList()
      _UpperCAmelCase						 =      nn.ParameterList()

      if div_val == 1:
        for i in range(len(self.cutoffs					)					):
          if d_proj != d_embed:
            self.out_projs.append(nn.Parameter(torch.FloatTensor(__lowerCAmelCase   ,				__lowerCAmelCase					)					)					)
          else:
            self.out_projs.append(__lowerCAmelCase					)

        self.out_layers.append(nn.Linear(__lowerCAmelCase   ,				__lowerCAmelCase					)					)
      else:
        for i in range(len(self.cutoffs					)					):
          _UpperCAmelCase ,				_UpperCAmelCase						 =      self.cutoff_ends[i], self.cutoff_ends[i + 1]
          _UpperCAmelCase						 =      d_embed // (div_val**i)

          self.out_projs.append(nn.Parameter(torch.FloatTensor(__lowerCAmelCase   ,				__lowerCAmelCase					)					)					)

          self.out_layers.append(nn.Linear(__lowerCAmelCase   ,				r_idx - l_idx					)					)

      _UpperCAmelCase						 =      keep_order


    def      lowerCAmelCase_  (					self	:  Optional[Any]   ,				__lowerCAmelCase	:  Optional[Any]   ,				__lowerCAmelCase	:  Dict   ,				__lowerCAmelCase	:  int   ,				__lowerCAmelCase	:  Tuple					):
      if proj is None:
        _UpperCAmelCase						 =      nn.functional.linear(__lowerCAmelCase   ,				__lowerCAmelCase   ,				bias=__lowerCAmelCase					)
      else:
        # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
        _UpperCAmelCase						 =      nn.functional.linear(__lowerCAmelCase   ,				proj.t().contiguous()					)
        _UpperCAmelCase						 =      nn.functional.linear(__lowerCAmelCase   ,				__lowerCAmelCase   ,				bias=__lowerCAmelCase					)
        # else:
        #     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
        #     if bias is not None:
        #         logit = logit + bias

      return logit


    def      lowerCAmelCase_  (					self	:  List[Any]   ,				__lowerCAmelCase	:  Optional[int]   ,				__lowerCAmelCase	:  Optional[Any]=None   ,				__lowerCAmelCase	:  Union[str, Any]=False					):

      if labels is not None:
        # Shift so that tokens < n predict n
        _UpperCAmelCase						 =      hidden[..., :-1, :].contiguous()
        _UpperCAmelCase						 =      labels[..., 1:].contiguous()
        _UpperCAmelCase						 =      hidden.view(-1   ,				hidden.size(-1					)					)
        _UpperCAmelCase						 =      labels.view(-1					)
        if hidden.size(0					) != labels.size(0					):
          raise RuntimeError("""Input and labels should have the same size in the batch dimension."""					)
      else:
        _UpperCAmelCase						 =      hidden.view(-1   ,				hidden.size(-1					)					)

      if self.n_clusters == 0:
        _UpperCAmelCase						 =      self._compute_logit(__lowerCAmelCase   ,				self.out_layers[0].weight   ,				self.out_layers[0].bias   ,				self.out_projs[0]					)
        if labels is not None:
          _UpperCAmelCase						 =      labels != -100
          _UpperCAmelCase						 =      torch.zeros_like(__lowerCAmelCase   ,				dtype=hidden.dtype   ,				device=hidden.device					)
          _UpperCAmelCase						 =      (
              -nn.functional.log_softmax(__lowerCAmelCase   ,				dim=-1					)[mask].gather(1   ,				labels[mask].unsqueeze(1					)					).squeeze(1					)
          )
        else:
          _UpperCAmelCase						 =      nn.functional.log_softmax(__lowerCAmelCase   ,				dim=-1					)
      else:
        # construct weights and biases
        _UpperCAmelCase ,				_UpperCAmelCase						 =      [], []
        for i in range(len(self.cutoffs					)					):
          if self.div_val == 1:
            _UpperCAmelCase ,				_UpperCAmelCase						 =      self.cutoff_ends[i], self.cutoff_ends[i + 1]
            _UpperCAmelCase						 =      self.out_layers[0].weight[l_idx:r_idx]
            _UpperCAmelCase						 =      self.out_layers[0].bias[l_idx:r_idx]
          else:
            _UpperCAmelCase						 =      self.out_layers[i].weight
            _UpperCAmelCase						 =      self.out_layers[i].bias

          if i == 0:
            _UpperCAmelCase						 =      torch.cat([weight_i, self.cluster_weight]   ,				dim=0					)
            _UpperCAmelCase						 =      torch.cat([bias_i, self.cluster_bias]   ,				dim=0					)

          weights.append(__lowerCAmelCase					)
          biases.append(__lowerCAmelCase					)

        _UpperCAmelCase ,				_UpperCAmelCase ,				_UpperCAmelCase						 =      weights[0], biases[0], self.out_projs[0]

        _UpperCAmelCase						 =      self._compute_logit(__lowerCAmelCase   ,				__lowerCAmelCase   ,				__lowerCAmelCase   ,				__lowerCAmelCase					)
        _UpperCAmelCase						 =      nn.functional.log_softmax(__lowerCAmelCase   ,				dim=1					)

        if labels is None:
          _UpperCAmelCase						 =      hidden.new_empty((head_logit.size(0					), self.n_token)					)
        else:
          _UpperCAmelCase						 =      torch.zeros_like(__lowerCAmelCase   ,				dtype=hidden.dtype   ,				device=hidden.device					)

        _UpperCAmelCase						 =      0
        _UpperCAmelCase						 =      [0] + self.cutoffs
        for i in range(len(__lowerCAmelCase					) - 1					):
          _UpperCAmelCase ,				_UpperCAmelCase						 =      cutoff_values[i], cutoff_values[i + 1]

          if labels is not None:
            _UpperCAmelCase						 =      (labels >= l_idx) & (labels < r_idx)
            _UpperCAmelCase						 =      mask_i.nonzero().squeeze()

            if indices_i.numel() == 0:
              continue

            _UpperCAmelCase						 =      labels.index_select(0   ,				__lowerCAmelCase					) - l_idx
            _UpperCAmelCase						 =      head_logprob.index_select(0   ,				__lowerCAmelCase					)
            _UpperCAmelCase						 =      hidden.index_select(0   ,				__lowerCAmelCase					)
          else:
            _UpperCAmelCase						 =      hidden

          if i == 0:
            if labels is not None:
              _UpperCAmelCase						 =      head_logprob_i.gather(1   ,				target_i[:, None]					).squeeze(1					)
            else:
              _UpperCAmelCase						 =      head_logprob[:, : self.cutoffs[0]]
          else:
            _UpperCAmelCase ,				_UpperCAmelCase ,				_UpperCAmelCase						 =      weights[i], biases[i], self.out_projs[i]

            _UpperCAmelCase						 =      self._compute_logit(__lowerCAmelCase   ,				__lowerCAmelCase   ,				__lowerCAmelCase   ,				__lowerCAmelCase					)
            _UpperCAmelCase						 =      nn.functional.log_softmax(__lowerCAmelCase   ,				dim=1					)
            _UpperCAmelCase						 =      self.cutoffs[0] + i - 1  # No probability for the head cluster
            if labels is not None:
              _UpperCAmelCase						 =      head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
                  1   ,				target_i[:, None]					).squeeze(1					)
            else:
              _UpperCAmelCase						 =      head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
              _UpperCAmelCase						 =      logprob_i

          if labels is not None:
            if (hasattr(self   ,				"""keep_order"""					) and self.keep_order) or keep_order:
              out.index_copy_(0   ,				__lowerCAmelCase   ,				-logprob_i					)
            else:
              out[offset : offset + logprob_i.size(0					)].copy_(-logprob_i					)
            offset += logprob_i.size(0					)

      return out


    def      lowerCAmelCase_  (					self	:  Union[str, Any]   ,				__lowerCAmelCase	:  List[Any]					):
      if self.n_clusters == 0:
        _UpperCAmelCase						 =      self._compute_logit(__lowerCAmelCase   ,				self.out_layers[0].weight   ,				self.out_layers[0].bias   ,				self.out_projs[0]					)
        return nn.functional.log_softmax(__lowerCAmelCase   ,				dim=-1					)
      else:
        # construct weights and biases
        _UpperCAmelCase ,				_UpperCAmelCase						 =      [], []
        for i in range(len(self.cutoffs					)					):
          if self.div_val == 1:
            _UpperCAmelCase ,				_UpperCAmelCase						 =      self.cutoff_ends[i], self.cutoff_ends[i + 1]
            _UpperCAmelCase						 =      self.out_layers[0].weight[l_idx:r_idx]
            _UpperCAmelCase						 =      self.out_layers[0].bias[l_idx:r_idx]
          else:
            _UpperCAmelCase						 =      self.out_layers[i].weight
            _UpperCAmelCase						 =      self.out_layers[i].bias

          if i == 0:
            _UpperCAmelCase						 =      torch.cat([weight_i, self.cluster_weight]   ,				dim=0					)
            _UpperCAmelCase						 =      torch.cat([bias_i, self.cluster_bias]   ,				dim=0					)

          weights.append(__lowerCAmelCase					)
          biases.append(__lowerCAmelCase					)

        _UpperCAmelCase ,				_UpperCAmelCase ,				_UpperCAmelCase						 =      weights[0], biases[0], self.out_projs[0]
        _UpperCAmelCase						 =      self._compute_logit(__lowerCAmelCase   ,				__lowerCAmelCase   ,				__lowerCAmelCase   ,				__lowerCAmelCase					)

        _UpperCAmelCase						 =      hidden.new_empty((head_logit.size(0					), self.n_token)					)
        _UpperCAmelCase						 =      nn.functional.log_softmax(__lowerCAmelCase   ,				dim=1					)

        _UpperCAmelCase						 =      [0] + self.cutoffs
        for i in range(len(__lowerCAmelCase					) - 1					):
          _UpperCAmelCase ,				_UpperCAmelCase						 =      cutoff_values[i], cutoff_values[i + 1]

          if i == 0:
            _UpperCAmelCase						 =      head_logprob[:, : self.cutoffs[0]]
          else:
            _UpperCAmelCase ,				_UpperCAmelCase ,				_UpperCAmelCase						 =      weights[i], biases[i], self.out_projs[i]

            _UpperCAmelCase						 =      self._compute_logit(__lowerCAmelCase   ,				__lowerCAmelCase   ,				__lowerCAmelCase   ,				__lowerCAmelCase					)
            _UpperCAmelCase						 =      nn.functional.log_softmax(__lowerCAmelCase   ,				dim=1					)

            _UpperCAmelCase						 =      head_logprob[:, -i] + tail_logprob_i
            _UpperCAmelCase						 =      logprob_i

        return out
