
'''simple docstring'''







import math
import os
import unittest

from transformers import MegatronBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import (
      MODEL_FOR_PRETRAINING_MAPPING,
      MegatronBertForCausalLM,
      MegatronBertForMaskedLM,
      MegatronBertForMultipleChoice,
      MegatronBertForNextSentencePrediction,
      MegatronBertForPreTraining,
      MegatronBertForQuestionAnswering,
      MegatronBertForSequenceClassification,
      MegatronBertForTokenClassification,
      MegatronBertModel,
  )







class UpperCAmelCase__	:





  """simple docstring"""

  def __init__(		self      :     Optional[Any]							,_a      :     List[str]							,_a      :     str=13							,_a      :     List[Any]=7							,_a      :     List[str]=True							,_a      :     Optional[int]=True							,_a      :     List[str]=True							,_a      :     Dict=True							,_a      :     Optional[Any]=99							,_a      :     Any=64							,_a      :     Optional[int]=32							,_a      :     List[str]=5							,_a      :     str=4							,_a      :     Union[str, Any]=37							,_a      :     Union[str, Any]="gelu"							,_a      :     Dict=0.1							,_a      :     Optional[Any]=0.1							,_a      :     Dict=512							,_a      :     Union[str, Any]=16							,_a      :     Optional[int]=2							,_a      :     str=0.02							,_a      :     Dict=3							,_a      :     str=4							,_a      :     Tuple=None							,):


     '''simple docstring'''






     _a				:	Optional[int]			=   parent
     _a				:	List[Any]			=   batch_size
     _a				:	Dict			=   seq_length
     _a				:	Any			=   is_training
     _a				:	Union[str, Any]			=   use_input_mask
     _a				:	Optional[int]			=   use_token_type_ids
     _a				:	Tuple			=   use_labels
     _a				:	Union[str, Any]			=   vocab_size
     _a				:	Dict			=   hidden_size
     _a				:	Union[str, Any]			=   embedding_size
     _a				:	List[Any]			=   num_hidden_layers
     _a				:	List[Any]			=   num_attention_heads
     _a				:	Union[str, Any]			=   intermediate_size
     _a				:	List[str]			=   hidden_act
     _a				:	Tuple			=   hidden_dropout_prob
     _a				:	str			=   attention_probs_dropout_prob
     _a				:	List[Any]			=   max_position_embeddings
     _a				:	List[str]			=   type_vocab_size
     _a				:	int			=   type_sequence_label_size
     _a				:	Union[str, Any]			=   initializer_range
     _a				:	List[Any]			=   num_labels
     _a				:	Any			=   num_choices
     _a				:	Optional[Any]			=   scope


  def   __lowercase		(		self      :     Optional[int]       ):


     '''simple docstring'''






     _a				:	List[str]			=   ids_tensor([self.batch_size, self.seq_length]							,self.vocab_size       )

     _a				:	List[Any]			=   None
     if self.use_input_mask:
        _a				:	List[Any]			=   random_attention_mask([self.batch_size, self.seq_length]       )

     _a				:	Optional[int]			=   None
     if self.use_token_type_ids:
        _a				:	int			=   ids_tensor([self.batch_size, self.seq_length]							,self.type_vocab_size       )

     _a				:	int			=   None
     _a				:	str			=   None
     _a				:	Optional[int]			=   None
     if self.use_labels:
        _a				:	Dict			=   ids_tensor([self.batch_size]							,self.type_sequence_label_size       )
        _a				:	int			=   ids_tensor([self.batch_size, self.seq_length]							,self.num_labels       )
        _a				:	str			=   ids_tensor([self.batch_size]							,self.num_choices       )

     _a				:	Dict			=   self.get_config()

     return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels


  def   __lowercase		(		self      :     List[Any]       ):


     '''simple docstring'''






     return MegatronBertConfig(
         vocab_size=self.vocab_size							,hidden_size=self.hidden_size							,num_hidden_layers=self.num_hidden_layers							,num_attention_heads=self.num_attention_heads							,intermediate_size=self.intermediate_size							,embedding_size=self.embedding_size							,hidden_act=self.hidden_act							,hidden_dropout_prob=self.hidden_dropout_prob							,attention_probs_dropout_prob=self.attention_probs_dropout_prob							,max_position_embeddings=self.max_position_embeddings							,type_vocab_size=self.type_vocab_size							,is_decoder=_a							,initializer_range=self.initializer_range							,)


  def   __lowercase		(		self      :     Tuple							,_a      :     int							,_a      :     str							,_a      :     str							,_a      :     Union[str, Any]							,_a      :     str							,_a      :     int							,_a      :     Any       ):


     '''simple docstring'''






     _a				:	Union[str, Any]			=   MegatronBertModel(config=_a       )
     model.to(_a       )
     model.eval()
     _a				:	str			=   model(_a							,attention_mask=_a							,token_type_ids=_a       )
     _a				:	str			=   model(_a							,token_type_ids=_a       )
     _a				:	int			=   model(_a       )

     self.parent.assertEqual(result.last_hidden_state.shape							,(self.batch_size, self.seq_length, self.hidden_size)       )
     self.parent.assertEqual(result.pooler_output.shape							,(self.batch_size, self.hidden_size)       )


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


     '''simple docstring'''






     _a				:	List[Any]			=   MegatronBertForMaskedLM(config=_a       )
     model.to(_a       )
     model.eval()
     _a				:	Dict			=   model(_a							,attention_mask=_a							,token_type_ids=_a							,labels=_a       )
     self.parent.assertEqual(result.logits.shape							,(self.batch_size, self.seq_length, self.vocab_size)       )


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


     '''simple docstring'''






     _a				:	Optional[Any]			=   MegatronBertForCausalLM(config=_a       )
     model.to(_a       )
     model.eval()
     _a				:	str			=   model(_a							,attention_mask=_a							,token_type_ids=_a							,labels=_a       )
     self.parent.assertEqual(result.logits.shape							,(self.batch_size, self.seq_length, self.vocab_size)       )


  def   __lowercase		(		self      :     Any							,_a      :     List[str]							,_a      :     List[str]							,_a      :     List[str]							,_a      :     List[str]							,_a      :     List[str]							,_a      :     List[Any]							,_a      :     Tuple       ):


     '''simple docstring'''






     _a				:	List[Any]			=   MegatronBertForNextSentencePrediction(config=_a       )
     model.to(_a       )
     model.eval()
     _a				:	List[str]			=   model(
         _a							,attention_mask=_a							,token_type_ids=_a							,labels=_a							,)
     self.parent.assertEqual(result.logits.shape							,(self.batch_size, 2)       )


  def   __lowercase		(		self      :     Tuple							,_a      :     List[str]							,_a      :     List[Any]							,_a      :     Any							,_a      :     List[str]							,_a      :     str							,_a      :     List[str]							,_a      :     Any       ):


     '''simple docstring'''






     _a				:	str			=   MegatronBertForPreTraining(config=_a       )
     model.to(_a       )
     model.eval()
     _a				:	Optional[Any]			=   model(
         _a							,attention_mask=_a							,token_type_ids=_a							,labels=_a							,next_sentence_label=_a							,)
     self.parent.assertEqual(result.prediction_logits.shape							,(self.batch_size, self.seq_length, self.vocab_size)       )
     self.parent.assertEqual(result.seq_relationship_logits.shape							,(self.batch_size, 2)       )


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


     '''simple docstring'''






     _a				:	int			=   MegatronBertForQuestionAnswering(config=_a       )
     model.to(_a       )
     model.eval()
     _a				:	List[str]			=   model(
         _a							,attention_mask=_a							,token_type_ids=_a							,start_positions=_a							,end_positions=_a							,)
     self.parent.assertEqual(result.start_logits.shape							,(self.batch_size, self.seq_length)       )
     self.parent.assertEqual(result.end_logits.shape							,(self.batch_size, self.seq_length)       )


  def   __lowercase		(		self      :     str							,_a      :     int							,_a      :     Dict							,_a      :     List[str]							,_a      :     Union[str, Any]							,_a      :     Tuple							,_a      :     Union[str, Any]							,_a      :     List[str]       ):


     '''simple docstring'''






     _a				:	Union[str, Any]			=   self.num_labels
     _a				:	List[Any]			=   MegatronBertForSequenceClassification(_a       )
     model.to(_a       )
     model.eval()
     _a				:	Tuple			=   model(_a							,attention_mask=_a							,token_type_ids=_a							,labels=_a       )
     self.parent.assertEqual(result.logits.shape							,(self.batch_size, self.num_labels)       )


  def   __lowercase		(		self      :     Tuple							,_a      :     Union[str, Any]							,_a      :     Dict							,_a      :     List[str]							,_a      :     Tuple							,_a      :     List[str]							,_a      :     Any							,_a      :     Optional[int]       ):


     '''simple docstring'''






     _a				:	str			=   self.num_labels
     _a				:	Any			=   MegatronBertForTokenClassification(config=_a       )
     model.to(_a       )
     model.eval()
     _a				:	List[Any]			=   model(_a							,attention_mask=_a							,token_type_ids=_a							,labels=_a       )
     self.parent.assertEqual(result.logits.shape							,(self.batch_size, self.seq_length, self.num_labels)       )


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


     '''simple docstring'''






     _a				:	Union[str, Any]			=   self.num_choices
     _a				:	str			=   MegatronBertForMultipleChoice(config=_a       )
     model.to(_a       )
     model.eval()
     _a				:	int			=   input_ids.unsqueeze(1       ).expand(-1							,self.num_choices							,-1       ).contiguous()
     _a				:	List[str]			=   token_type_ids.unsqueeze(1       ).expand(-1							,self.num_choices							,-1       ).contiguous()
     _a				:	Dict			=   input_mask.unsqueeze(1       ).expand(-1							,self.num_choices							,-1       ).contiguous()
     _a				:	Optional[int]			=   model(
         _a							,attention_mask=_a							,token_type_ids=_a							,labels=_a							,)
     self.parent.assertEqual(result.logits.shape							,(self.batch_size, self.num_choices)       )


  def   __lowercase		(		self      :     Dict       ):


     '''simple docstring'''






     _a				:	Dict			=   self.prepare_config_and_inputs()
     (
         (
         _a
     ),  (
         _a
     ),  (
         _a
     ),  (
         _a
     ),  (
         _a
     ),  (
         _a
     ),  (
         _a
     ),  
     )				:	Union[str, Any]			=   config_and_inputs
     _a				:	Union[str, Any]			=   {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
     return config, inputs_dict







@require_torch
class UpperCAmelCase__	( lowercase__    ,				lowercase__    ,				unittest.TestCase	):





  """simple docstring"""

  __UpperCAmelCase				:				str   =					(
      (
          MegatronBertModel,
          MegatronBertForMaskedLM,
          MegatronBertForCausalLM,
          MegatronBertForMultipleChoice,
          MegatronBertForNextSentencePrediction,
          MegatronBertForPreTraining,
          MegatronBertForQuestionAnswering,
          MegatronBertForSequenceClassification,
          MegatronBertForTokenClassification,
      )
      if is_torch_available()
      else ()
  )
  __UpperCAmelCase				:				Tuple   =					(
      {
          '''feature-extraction''': MegatronBertModel,
          '''fill-mask''': MegatronBertForMaskedLM,
          '''question-answering''': MegatronBertForQuestionAnswering,
          '''text-classification''': MegatronBertForSequenceClassification,
          '''text-generation''': MegatronBertForCausalLM,
          '''token-classification''': MegatronBertForTokenClassification,
          '''zero-shot''': MegatronBertForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )
  __UpperCAmelCase				:				List[Any]   =					True
  # test_resize_embeddings = False
  __UpperCAmelCase				:				List[Any]   =					False


  def   __lowercase		(		self      :     Any							,_a      :     List[Any]							,_a      :     List[str]							,_a      :     List[str]=False       ):


     '''simple docstring'''






     _a				:	Dict			=   super()._prepare_for_class(_a							,_a							,return_labels=_a       )

     if return_labels:
        if model_class in get_values(_a       ):
           _a				:	Dict			=   torch.zeros(
               (self.model_tester.batch_size, self.model_tester.seq_length)							,dtype=torch.long							,device=_a       )
           _a				:	Dict			=   torch.zeros(
               self.model_tester.batch_size							,dtype=torch.long							,device=_a       )
     return inputs_dict


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


     '''simple docstring'''






     _a				:	Dict			=   MegatronBertModelTester(self       )
     _a				:	Any			=   ConfigTester(self							,config_class=_a							,hidden_size=37       )


  def   __lowercase		(		self      :     str       ):


     '''simple docstring'''






     self.config_tester.run_common_tests()


  def   __lowercase		(		self      :     Optional[Any]       ):


     '''simple docstring'''






     _a				:	List[Any]			=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_model(*_a       )


  def   __lowercase		(		self      :     Optional[Any]       ):


     '''simple docstring'''






     _a				:	List[str]			=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_masked_lm(*_a       )


  def   __lowercase		(		self      :     int       ):


     '''simple docstring'''






     _a				:	Any			=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*_a       )


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


     '''simple docstring'''






     _a				:	str			=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*_a       )


  def   __lowercase		(		self      :     List[Any]       ):


     '''simple docstring'''






     _a				:	Tuple			=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_pretraining(*_a       )


  def   __lowercase		(		self      :     Any       ):


     '''simple docstring'''






     _a				:	Tuple			=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_question_answering(*_a       )


  def   __lowercase		(		self      :     Any       ):


     '''simple docstring'''






     _a				:	Union[str, Any]			=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*_a       )


  def   __lowercase		(		self      :     List[str]       ):


     '''simple docstring'''






     _a				:	List[str]			=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_token_classification(*_a       )





def 					UpperCAmelCase_		(__a      :    Tuple ):






   """simple docstring"""

   return torch.tensor(
       __a ,	dtype=torch.long ,	device=__a ,	)


__lowerCAmelCase         =    1e-4







@require_torch
@require_sentencepiece
@require_tokenizers
class UpperCAmelCase__	( unittest.TestCase	):





  """simple docstring"""

  @slow
  @unittest.skip('Model is not available.'       )
  def   __lowercase		(		self      :     str       ):


     '''simple docstring'''






     _a				:	int			=   'nvidia/megatron-bert-uncased-345m'
     if "MYDIR" in os.environ:
        _a				:	Optional[int]			=   os.path.join(os.environ['MYDIR']							,_a       )
     _a				:	List[Any]			=   MegatronBertModel.from_pretrained(_a       )
     model.to(_a       )
     model.half()
     _a				:	List[str]			=   _long_tensor([[101, 7110, 1005, 1056, 2023, 1_1333, 1_7413, 1029, 102]]       )
     with torch.no_grad():
        _a				:	str			=   model(_a       )[0]
     _a				:	Tuple			=   torch.Size((1, 9, 1024)       )
     self.assertEqual(output.shape							,_a       )

     _a				:	str			=   [-0.6040, -0.2517, -0.1025, 0.3420, -0.6758, -0.0017, -0.1089, -0.1990, 0.5728]
     for ii in range(3       ):
        for jj in range(3       ):
           _a				:	Optional[Any]			=   output[0, ii, jj]
           _a				:	Optional[int]			=   expected[3 * ii + jj]
           _a				:	Optional[Any]			=   'ii={} jj={} a={} b={}'.format(_a							,_a							,_a							,_a       )
           self.assertTrue(math.isclose(_a							,_a							,rel_tol=_a							,abs_tol=_a       )							,msg=_a       )





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



