




import json
import logging
import math
import os
import sys
from dataclasses import dataclass, field
from typing import Optional

from datasets import Dataset, load_dataset

import transformers
from transformers import (
    CONFIG_MAPPING,
    MODEL_FOR_MASKED_LM_MAPPING,
    AutoConfig,
    AutoModelForMaskedLM,
    AutoTokenizer,
    DataCollatorForWholeWordMask,
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    set_seed,
)
from transformers.trainer_utils import get_last_checkpoint, is_main_process


UpperCAmelCase_									=	logging.getLogger(__name__)
UpperCAmelCase_									=	list(MODEL_FOR_MASKED_LM_MAPPING.keys())
UpperCAmelCase_									=	tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)





@dataclass
class     lowerCamelCase__:
      UpperCAmelCase__      :	Optional[str]    		=   field(
          default=__lowerCamelCase ,     metadata={
              'help': (
                  'The model checkpoint for weights initialization.Don\'t set if you want to train a model from scratch.'
              )
          } ,     )
      UpperCAmelCase__      :	Optional[str]    		=   field(
          default=__lowerCamelCase ,     metadata={'help': 'If training from scratch, pass a model type from the list: ' + ', '.join(__lowerCamelCase)} ,     )
      UpperCAmelCase__      :	Optional[str]    		=   field(
          default=__lowerCamelCase ,     metadata={
              'help': (
                  'Override some existing default config settings when a model is trained from scratch. Example: '
                  'n_embd=10,resid_pdrop=0.2,scale_attn_weights=false,summary_type=cls_index'
              )
          } ,     )
      UpperCAmelCase__      :	Optional[str]    		=   field(
          default=__lowerCamelCase ,     metadata={'help': 'Pretrained config name or path if not the same as model_name'})
      UpperCAmelCase__      :	Optional[str]    		=   field(
          default=__lowerCamelCase ,     metadata={'help': 'Pretrained tokenizer name or path if not the same as model_name'})
      UpperCAmelCase__      :	Optional[str]    		=   field(
          default=__lowerCamelCase ,     metadata={'help': 'Where do you want to store the pretrained models downloaded from huggingface.co'} ,     )
      UpperCAmelCase__      :	bool    		=   field(
          default=__lowerCamelCase ,     metadata={'help': 'Whether to use one of the fast tokenizer (backed by the tokenizers library) or not.'} ,     )
      UpperCAmelCase__      :	str    		=   field(
          default='main' ,     metadata={'help': 'The specific model version to use (can be a branch name, tag name or commit id).'} ,     )
      UpperCAmelCase__      :	bool    		=   field(
          default=__lowerCamelCase ,     metadata={
              'help': (
                  'Will use the token generated when running `huggingface-cli login` (necessary to use this script '
                  'with private models).'
              )
          } ,     )
      def 			lowerCAmelCase__						(				self: List[Any]   ):
             if self.config_overrides is not None and (self.config_name is not None or self.model_name_or_path is not None):
                    raise ValueError(
                        """--config_overrides can't be used in combination with --config_name or --model_name_or_path"""   )





@dataclass
class     lowerCamelCase__:
      UpperCAmelCase__      :	Optional[str]    		=   field(
          default=__lowerCamelCase ,     metadata={'help': 'The name of the dataset to use (via the datasets library).'})
      UpperCAmelCase__      :	Optional[str]    		=   field(
          default=__lowerCamelCase ,     metadata={'help': 'The configuration name of the dataset to use (via the datasets library).'})
      UpperCAmelCase__      :	Optional[str]    		=   field(default=__lowerCamelCase ,     metadata={'help': 'The input training data file (a text file).'})
      UpperCAmelCase__      :	Optional[str]    		=   field(
          default=__lowerCamelCase ,     metadata={'help': 'An optional input evaluation data file to evaluate the perplexity on (a text file).'} ,     )
      UpperCAmelCase__      :	Optional[str]    		=   field(
          default=__lowerCamelCase ,     metadata={'help': 'An optional input train ref data file for whole word masking in Chinese.'} ,     )
      UpperCAmelCase__      :	Optional[str]    		=   field(
          default=__lowerCamelCase ,     metadata={'help': 'An optional input validation ref data file for whole word masking in Chinese.'} ,     )
      UpperCAmelCase__      :	bool    		=   field(
          default=__lowerCamelCase ,     metadata={'help': 'Overwrite the cached training and evaluation sets'})
      UpperCAmelCase__      :	Optional[int]    		=   field(
          default=5 ,     metadata={
              'help': 'The percentage of the train set used as validation set in case there\'s no validation split'
          } ,     )
      UpperCAmelCase__      :	Optional[int]    		=   field(
          default=__lowerCamelCase ,     metadata={
              'help': (
                  'The maximum total input sequence length after tokenization. Sequences longer '
                  'than this will be truncated. Default to the max input length of the model.'
              )
          } ,     )
      UpperCAmelCase__      :	Optional[int]    		=   field(
          default=__lowerCamelCase ,     metadata={'help': 'The number of processes to use for the preprocessing.'} ,     )
      UpperCAmelCase__      :	float    		=   field(
          default=0.15 ,     metadata={'help': 'Ratio of tokens to mask for masked language modeling loss'})
      UpperCAmelCase__      :	bool    		=   field(
          default=__lowerCamelCase ,     metadata={
              'help': (
                  'Whether to pad all samples to `max_seq_length`. '
                  'If False, will pad the samples dynamically when batching to the maximum length in the batch.'
              )
          } ,     )
      def 			lowerCAmelCase__						(				self: str   ):
             if self.train_file is not None:
                    __lowerCamelCase   =						self.train_file.split("""."""   )[-1]
                    assert extension in ["csv", "json", "txt"], "`train_file` should be a csv, a json or a txt file."
             if self.validation_file is not None:
                    __lowerCamelCase   =						self.validation_file.split("""."""   )[-1]
                    assert extension in ["csv", "json", "txt"], "`validation_file` should be a csv, a json or a txt file."




def 					lowerCamelCase__       (						A__       :					List[str]  ,			A__       :					List[str]	):





       '''simple docstring'''



       with open(A__  ,			"""r"""  ,			encoding="""utf-8"""	) as f:
              __lowerCamelCase   =						[json.loads(A__	) for line in f.read().splitlines() if (len(A__	) > 0 and not line.isspace())]
       assert len(A__	) == len(A__	)

       __lowerCamelCase   =						{c: dataset[c] for c in dataset.column_names}
       __lowerCamelCase   =						refs
       return Dataset.from_dict(A__	)




def 					lowerCamelCase__       (						):





       '''simple docstring'''



       __lowerCamelCase   =						HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)	)
       if len(sys.argv	) == 2 and sys.argv[1].endswith(""".json"""	):
              # If we pass only one argument to the script and it's the path to a json file,
              # let's parse it to get our arguments.
              __lowerCamelCase,   __lowerCamelCase,   __lowerCamelCase   =						parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]	)	)
       else:
              __lowerCamelCase,   __lowerCamelCase,   __lowerCamelCase   =						parser.parse_args_into_dataclasses()

       # Detecting last checkpoint.
       __lowerCamelCase   =						None
       if os.path.isdir(training_args.output_dir	) and training_args.do_train and not training_args.overwrite_output_dir:
              __lowerCamelCase   =						get_last_checkpoint(training_args.output_dir	)
              if last_checkpoint is None and len(os.listdir(training_args.output_dir	)	) > 0:
                     raise ValueError(
                         f'Output directory ({training_args.output_dir}) already exists and is not empty. '
                         """Use --overwrite_output_dir to overcome."""	)
              elif last_checkpoint is not None:
                     logger.info(
                         f'Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change '
                         """the `--output_dir` or add `--overwrite_output_dir` to train from scratch."""	)

    # Setup logging
       logging.basicConfig(
           format="""%(asctime)s - %(levelname)s - %(name)s - %(message)s"""  ,			datefmt="""%m/%d/%Y %H:%M:%S"""  ,			handlers=[logging.StreamHandler(sys.stdout	)]  ,			)
       logger.setLevel(logging.INFO if is_main_process(training_args.local_rank	) else logging.WARN	)

       # Log on each process the small summary:
       logger.warning(
           f'Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}'
           + f'distributed training: {bool(training_args.local_rank != -1	)}, 16-bits training: {training_args.fpaa}'	)
       # Set the verbosity to info of the Transformers logger (on main process only):
       if is_main_process(training_args.local_rank	):
              transformers.utils.logging.set_verbosity_info()
              transformers.utils.logging.enable_default_handler()
              transformers.utils.logging.enable_explicit_format()
       logger.info("""Training/evaluation parameters %s"""  ,			A__	)

       # Set seed before initializing model.
       set_seed(training_args.seed	)

       # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
       # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
       # (the dataset will be downloaded automatically from the datasets Hub).
       #
       # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called
       # 'text' is found. You can easily tweak this behavior (see below).
       #
       # In distributed training, the load_dataset function guarantee that only one local process can concurrently
       # download the dataset.
       if data_args.dataset_name is not None:
              # Downloading and loading a dataset from the hub.
              __lowerCamelCase   =						load_dataset(data_args.dataset_name  ,			data_args.dataset_config_name	)
              if "validation" not in datasets.keys():
                     __lowerCamelCase   =						load_dataset(
                         data_args.dataset_name  ,			data_args.dataset_config_name  ,			split=f'train[:{data_args.validation_split_percentage}%]'  ,			)
                     __lowerCamelCase   =						load_dataset(
                         data_args.dataset_name  ,			data_args.dataset_config_name  ,			split=f'train[{data_args.validation_split_percentage}%:]'  ,			)
       else:
              __lowerCamelCase   =						{}
              if data_args.train_file is not None:
                     __lowerCamelCase   =						data_args.train_file
              if data_args.validation_file is not None:
                     __lowerCamelCase   =						data_args.validation_file
              __lowerCamelCase   =						data_args.train_file.split("""."""	)[-1]
              if extension == "txt":
                     __lowerCamelCase   =						"""text"""
              __lowerCamelCase   =						load_dataset(A__  ,			data_files=A__	)
       # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
       # https://huggingface.co/docs/datasets/loading_datasets.html.

       # Load pretrained model and tokenizer
       #
       # Distributed training:
       # The .from_pretrained methods guarantee that only one local process can concurrently
       # download model & vocab.
       __lowerCamelCase   =						{
           """cache_dir""": model_args.cache_dir,
           """revision""": model_args.model_revision,
           """use_auth_token""": True if model_args.use_auth_token else None,
       }
       if model_args.config_name:
              __lowerCamelCase   =						AutoConfig.from_pretrained(model_args.config_name  ,			**A__	)
       elif model_args.model_name_or_path:
              __lowerCamelCase   =						AutoConfig.from_pretrained(model_args.model_name_or_path  ,			**A__	)
       else:
              __lowerCamelCase   =						CONFIG_MAPPING[model_args.model_type]()
              logger.warning("""You are instantiating a new config instance from scratch."""	)
              if model_args.config_overrides is not None:
                     logger.info(f'Overriding config: {model_args.config_overrides}'	)
                     config.update_from_string(model_args.config_overrides	)
                     logger.info(f'New config: {config}'	)

       __lowerCamelCase   =						{
           """cache_dir""": model_args.cache_dir,
           """use_fast""": model_args.use_fast_tokenizer,
           """revision""": model_args.model_revision,
           """use_auth_token""": True if model_args.use_auth_token else None,
       }
       if model_args.tokenizer_name:
              __lowerCamelCase   =						AutoTokenizer.from_pretrained(model_args.tokenizer_name  ,			**A__	)
       elif model_args.model_name_or_path:
              __lowerCamelCase   =						AutoTokenizer.from_pretrained(model_args.model_name_or_path  ,			**A__	)
       else:
              raise ValueError(
                  """You are instantiating a new tokenizer from scratch. This is not supported by this script."""
                  """You can do it from another script, save it, and load it from here, using --tokenizer_name."""	)

       if model_args.model_name_or_path:
              __lowerCamelCase   =						AutoModelForMaskedLM.from_pretrained(
                  model_args.model_name_or_path  ,			from_tf=bool(""".ckpt""" in model_args.model_name_or_path	)  ,			config=A__  ,			cache_dir=model_args.cache_dir  ,			revision=model_args.model_revision  ,			use_auth_token=True if model_args.use_auth_token else None  ,			)
       else:
              logger.info("""Training new model from scratch"""	)
              __lowerCamelCase   =						AutoModelForMaskedLM.from_config(A__	)

       model.resize_token_embeddings(len(A__	)	)

       # Preprocessing the datasets.
       # First we tokenize all the texts.
       if training_args.do_train:
              __lowerCamelCase   =						datasets["""train"""].column_names
       else:
              __lowerCamelCase   =						datasets["""validation"""].column_names
       __lowerCamelCase   =						"""text""" if """text""" in column_names else column_names[0]

       __lowerCamelCase   =						"""max_length""" if data_args.pad_to_max_length else False

       def tokenize_function(A__       :					Optional[Any]	):
              # Remove empty lines
              __lowerCamelCase   =						[line for line in examples["""text"""] if len(A__	) > 0 and not line.isspace()]
              return tokenizer(examples["""text"""]  ,			padding=A__  ,			truncation=A__  ,			max_length=data_args.max_seq_length	)

       __lowerCamelCase   =						datasets.map(
           A__  ,			batched=A__  ,			num_proc=data_args.preprocessing_num_workers  ,			remove_columns=[text_column_name]  ,			load_from_cache_file=not data_args.overwrite_cache  ,			)

       # Add the chinese references if provided
       if data_args.train_ref_file is not None:
              __lowerCamelCase   =						add_chinese_references(tokenized_datasets["""train"""]  ,			data_args.train_ref_file	)
       if data_args.validation_ref_file is not None:
              __lowerCamelCase   =						add_chinese_references(
                  tokenized_datasets["""validation"""]  ,			data_args.validation_ref_file	)
       # If we have ref files, need to avoid it removed by trainer
       __lowerCamelCase   =						data_args.train_ref_file or data_args.validation_ref_file
       if has_ref:
              __lowerCamelCase   =						False

       # Data collator
       # This one will take care of randomly masking the tokens.
       __lowerCamelCase   =						DataCollatorForWholeWordMask(tokenizer=A__  ,			mlm_probability=data_args.mlm_probability	)

       # Initialize our Trainer
       __lowerCamelCase   =						Trainer(
           model=A__  ,			args=A__  ,			train_dataset=tokenized_datasets["""train"""] if training_args.do_train else None  ,			eval_dataset=tokenized_datasets["""validation"""] if training_args.do_eval else None  ,			tokenizer=A__  ,			data_collator=A__  ,			)

       # Training
       if training_args.do_train:
              if last_checkpoint is not None:
                     __lowerCamelCase   =						last_checkpoint
              elif model_args.model_name_or_path is not None and os.path.isdir(model_args.model_name_or_path	):
                     __lowerCamelCase   =						model_args.model_name_or_path
              else:
                     __lowerCamelCase   =						None
              __lowerCamelCase   =						trainer.train(resume_from_checkpoint=A__	)
              trainer.save_model()  # Saves the tokenizer too for easy upload

              __lowerCamelCase   =						os.path.join(training_args.output_dir  ,			"""train_results.txt"""	)
              if trainer.is_world_process_zero():
                     with open(A__  ,			"""w"""	) as writer:
                            logger.info("""***** Train results *****"""	)
                            for key, value in sorted(train_result.metrics.items()	):
                                   logger.info(f'  {key} = {value}'	)
                                   writer.write(f'{key} = {value}\n'	)

            # Need to save the state, since Trainer.save_model saves only the tokenizer with the model
                     trainer.state.save_to_json(os.path.join(training_args.output_dir  ,			"""trainer_state.json"""	)	)

    # Evaluation
       __lowerCamelCase   =						{}
       if training_args.do_eval:
              logger.info("""*** Evaluate ***"""	)

              __lowerCamelCase   =						trainer.evaluate()

              __lowerCamelCase   =						math.exp(eval_output["""eval_loss"""]	)
              __lowerCamelCase   =						perplexity

              __lowerCamelCase   =						os.path.join(training_args.output_dir  ,			"""eval_results_mlm_wwm.txt"""	)
              if trainer.is_world_process_zero():
                     with open(A__  ,			"""w"""	) as writer:
                            logger.info("""***** Eval results *****"""	)
                            for key, value in sorted(results.items()	):
                                   logger.info(f'  {key} = {value}'	)
                                   writer.write(f'{key} = {value}\n'	)

       return results




def 					lowerCamelCase__       (						A__       :					Dict	):





       '''simple docstring'''



       main()


if __name__ == "__main__":
       main()



import json
import multiprocessing
import os
import re
from collections import defaultdict

import torch
from accelerate import Accelerator
from accelerate.utils import set_seed
from arguments import HumanEvalArguments
from datasets import load_dataset, load_metric
from torch.utils.data import IterableDataset
from torch.utils.data.dataloader import DataLoader
from tqdm import tqdm

import transformers
from transformers import AutoModelForCausalLM, AutoTokenizer, HfArgumentParser, StoppingCriteria, StoppingCriteriaList


__lowerCAmelCase			:     Optional[int]           =   ['\nclass', '\ndef', '\n#', '\n@', '\nprint', '\nif']




class        UpperCAmelCase_			( _A     ):


  '''simple docstring'''







  def __init__( self			: Union[str, Any]			,							UpperCamelCase__			: List[Any]			,							UpperCamelCase__			: Optional[int]			,							UpperCamelCase__			: Any=None			,							UpperCamelCase__			: Union[str, Any]=1    )			->	str:







    """simple docstring"""






    __magic_name__				     =       tokenizer
    __magic_name__				     =       dataset
    __magic_name__				     =       len(UpperCamelCase__    ) if n_tasks is None else n_tasks
    __magic_name__				     =       n_copies





  def __iter__( self			: List[Any]    )			->	List[str]:







    """simple docstring"""






    __magic_name__				     =       []
    for task in range(self.n_tasks    ):
      # without strip, the model generate commented codes ...
      prompts.append(self.tokenizer.eos_token + self.dataset[task]["""prompt"""].strip()    )
    __magic_name__				     =       self.tokenizer(UpperCamelCase__			,							padding=UpperCamelCase__			,							return_tensors="""pt"""    )
    for task in range(self.n_tasks    ):
      for _ in range(self.n_copies    ):
        yield {
            "ids": outputs.input_ids[task],
            "task_id": task,
            "input_len": outputs.attention_mask[task].sum(),
        }







class        UpperCAmelCase_			( _A     ):


  '''simple docstring'''







  def __init__( self			: Union[str, Any]			,							UpperCamelCase__			: Dict			,							UpperCamelCase__			: str			,							UpperCamelCase__			: str    )			->	List[str]:







    """simple docstring"""






    __magic_name__				     =       start_length
    __magic_name__				     =       eof_strings
    __magic_name__				     =       tokenizer





  def __call__( self			: List[Any]			,							UpperCamelCase__			: Tuple			,							UpperCamelCase__			: Optional[Any]			,							**UpperCamelCase__			: Optional[int]    )			->	List[str]:







    """simple docstring"""






    __magic_name__				     =       self.tokenizer.batch_decode(input_ids[:, self.start_length :]    )
    __magic_name__				     =       []
    for decoded_generation in decoded_generations:
      done.append(any(stop_string in decoded_generation for stop_string in self.eof_strings    )    )
    return all(UpperCamelCase__    )







def a__    (       A_					):






  '''simple docstring'''





  __magic_name__				     =       re.split("""(%s)""" % """|""".join(A_					),     A_					)
  # last string should be ""
  return "".join(string_list[:-2]					)

def a__    (       A_,     A_,     A_,     A_,     A_,     A_=20,     **A_					):






  '''simple docstring'''





  __magic_name__				     =       defaultdict(A_					)  # dict of list of generated tokens
  for step, batch in tqdm(enumerate(A_					)					):
    with torch.no_grad():
      __magic_name__				     =       batch["""ids"""].shape[-1]
      __magic_name__				     =       accelerator.unwrap_model(A_					).generate(
          input_ids=batch["""ids"""][:, : batch["""input_len"""]],     num_return_sequences=A_,     **A_					)
      # each task is generated batch_size times
      __magic_name__				     =       batch["""task_id"""].repeat(A_					)
      __magic_name__				     =       accelerator.pad_across_processes(
          A_,     dim=1,     pad_index=tokenizer.pad_token_id					)

      __magic_name__     ,    __magic_name__				     =       accelerator.gather((generated_tokens, generated_tasks)					)
      __magic_name__				     =       generated_tokens.cpu().numpy()
      __magic_name__				     =       generated_tasks.cpu().numpy()

      for task, generated_tokens in zip(A_,     A_					):
        gen_token_dict[task].append(A_					)

  __magic_name__				     =       [[] for _ in range(A_					)]
  for task, generated_tokens in gen_token_dict.items():
    for s in generated_tokens:
      __magic_name__				     =       tokenizer.decode(A_,     skip_special_tokens=A_,     clean_up_tokenization_spaces=A_					)
      code_gens[task].append(remove_last_block(A_					)					)
  return code_gens

def a__    (       ):






  '''simple docstring'''





  __magic_name__				     =       HfArgumentParser(A_					)
  __magic_name__				     =       parser.parse_args()

  transformers.logging.set_verbosity_error()
  # enables code execution in code_eval metric
  __magic_name__				     =       args.HF_ALLOW_CODE_EVAL
  # make sure tokenizer plays nice with multiprocessing
  __magic_name__				     =       """false"""

  if args.num_workers is None:
    __magic_name__				     =       multiprocessing.cpu_count()

  # Use dataset load to feed to accelerate
  __magic_name__				     =       Accelerator()
  set_seed(args.seed,     device_specific=A_					)

  # Load model and tokenizer
  __magic_name__				     =       AutoTokenizer.from_pretrained(args.model_ckpt					)
  __magic_name__				     =       tokenizer.eos_token
  __magic_name__				     =       AutoModelForCausalLM.from_pretrained(args.model_ckpt					)

  # Generation settings
  __magic_name__				     =       {
      """do_sample""": args.do_sample,
      """temperature""": args.temperature,
      """max_new_tokens""": args.max_new_tokens,
      """top_p""": args.top_p,
      """top_k""": args.top_k,
      """stopping_criteria""": StoppingCriteriaList([EndOfFunctionCriteria(0,     A_,     A_					)]					),
  }

  # Load evaluation dataset and metric
  __magic_name__				     =       load_dataset("""openai_humaneval"""					)
  __magic_name__				     =       load_metric("""code_eval"""					)

  __magic_name__				     =       args.num_tasks if args.num_tasks is not None else len(human_eval["""test"""]					)
  __magic_name__				     =       args.n_samples // args.batch_size

  __magic_name__				     =       TokenizedDataset(A_,     human_eval["""test"""],     n_copies=A_,     n_tasks=A_					)
  # do not confuse args.batch_size, which is actually the num_return_sequences
  __magic_name__				     =       DataLoader(A_,     batch_size=1					)

  # Run a quick test to see if code evaluation is enabled
  try:
    __magic_name__				     =       code_eval_metric.compute(references=[""""""],     predictions=[[""""""]]					)
  except ValueError as exception:
    print(
        """Code evaluation not enabled. Read the warning below carefully and then use `--HF_ALLOW_CODE_EVAL=\"1\"`"""
        """ flag to enable code evaluation."""					)
    raise exception

  __magic_name__     ,    __magic_name__				     =       accelerator.prepare(A_,     A_					)

  __magic_name__				     =       complete_code(
      A_,     A_,     A_,     A_,     n_tasks=A_,     batch_size=args.batch_size,     **A_,     )

  if accelerator.is_main_process:
    __magic_name__				     =       []

    for task in tqdm(range(A_					)					):
      __magic_name__				     =       human_eval["""test"""][task]["""test"""]
      __magic_name__				     =       f'''check({human_eval['test'][task]['entry_point']})'''
      references.append("""\n""" + test_func + """\n""" + entry_point					)

    # Evaluate completions with "code_eval" metric
    __magic_name__     ,    __magic_name__				     =       code_eval_metric.compute(
        references=A_,     predictions=A_,     num_workers=args.num_workers					)
    print(f'''Results: {pass_at_k}'''					)

    # Save results to json file
    with open(args.output_file,     """w"""					) as fp:
      json.dump(A_,     A_					)


# For some reason the folliwng seems to be necessary sometimes for code_eval to work nice with multiprocessing
# https://stackoverflow.com/questions/60804599/python-multiprocessing-keeps-spawning-the-whole-script
if __name__ == "__main__":
    main()

