'''simple docstring'''




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


__UpperCAmelCase												=["\nclass", "\ndef", "\n#", "\n@", "\nprint", "\nif"]




class    a__				(       UpperCAmelCase__			):






     def __init__(			self		:		Dict  ,    a		:		List[str]  ,    a		:		Optional[int]  ,    a		:		List[Any]=None  ,    a		:		Tuple=1   ):





        """simple docstring"""



        __lowerCamelCase          =  tokenizer
        __lowerCamelCase          =  dataset
        __lowerCamelCase          =  len(a   ) if n_tasks is None else n_tasks
        __lowerCamelCase          =  n_copies






     def __iter__(			self		:		int   ):





        """simple docstring"""



        __lowerCamelCase          =  []
        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()   )
        __lowerCamelCase          =  self.tokenizer(a  ,    padding=a  ,    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    a__				(       UpperCAmelCase__			):






     def __init__(			self		:		Dict  ,    a		:		str  ,    a		:		int  ,    a		:		Optional[Any]   ):





        """simple docstring"""



        __lowerCamelCase          =  start_length
        __lowerCamelCase          =  eof_strings
        __lowerCamelCase          =  tokenizer






     def __call__(			self		:		Optional[int]  ,    a		:		Optional[int]  ,    a		:		Any  ,    **a		:		str   ):





        """simple docstring"""



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






def __lowerCAmelCase			(							UpperCamelCase__    )     ->   Any:
   __lowerCamelCase          =  re.split('''(%s)''' % '''|'''.join(UpperCamelCase__    )    ,  UpperCamelCase__    )
   # last string should be ""
   return "".join(string_list[:-2]    )






def __lowerCAmelCase			(							UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__=20    ,  **UpperCamelCase__    )     ->   List[Any]:
   __lowerCamelCase          =  defaultdict(UpperCamelCase__    )  # dict of list of generated tokens
   for step, batch in tqdm(enumerate(UpperCamelCase__    )    ):
      with torch.no_grad():
         __lowerCamelCase          =  batch['''ids'''].shape[-1]
         __lowerCamelCase          =  accelerator.unwrap_model(UpperCamelCase__    ).generate(
             input_ids=batch['''ids'''][:, : batch['''input_len''']]    ,  num_return_sequences=UpperCamelCase__    ,  **UpperCamelCase__    )
         # each task is generated batch_size times
         __lowerCamelCase          =  batch['''task_id'''].repeat(UpperCamelCase__    )
         __lowerCamelCase          =  accelerator.pad_across_processes(
             UpperCamelCase__    ,  dim=1    ,  pad_index=tokenizer.pad_token_id    )

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

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

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






def __lowerCAmelCase			(							)     ->   Optional[Any]:
   # Setup configuration
   __lowerCamelCase          =  HfArgumentParser(UpperCamelCase__    )
   __lowerCamelCase          =  parser.parse_args()

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

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

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

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

   # Generation settings
   __lowerCamelCase          =  {
       '''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    ,  UpperCamelCase__    ,  UpperCamelCase__    )]    ),
   }

   # Load evaluation dataset and metric
   __lowerCamelCase          =  load_dataset('''openai_humaneval'''    )
   __lowerCamelCase          =  load_metric('''code_eval'''    )

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

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

   # Run a quick test to see if code evaluation is enabled
   try:
      __lowerCamelCase          =  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

   __lowerCamelCase     ,  __lowerCamelCase          =  accelerator.prepare(UpperCamelCase__    ,  UpperCamelCase__    )

   __lowerCamelCase          =  complete_code(
       UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  n_tasks=UpperCamelCase__    ,  batch_size=args.batch_size    ,  **UpperCamelCase__    ,  )

   if accelerator.is_main_process:
      __lowerCamelCase          =  []

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

      # Evaluate completions with "code_eval" metric
      __lowerCamelCase     ,  __lowerCamelCase          =  code_eval_metric.compute(
          references=UpperCamelCase__    ,  predictions=UpperCamelCase__    ,  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(UpperCamelCase__    ,  UpperCamelCase__    )


# 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()

'''simple docstring'''




from __future__ import annotations






def __lowerCAmelCase			(							UpperCamelCase__    ,  UpperCamelCase__    )     ->   list[list[int]]:
    __lowerCamelCase          =  []
    create_all_state(1    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  []    ,  UpperCamelCase__    )
    return result






def __lowerCAmelCase			(							UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  )     ->   None:
    if level == 0:
        total_list.append(current_list[:]    )
        return

    for i in range(UpperCamelCase__    ,  total_number - level + 2    ):
        current_list.append(UpperCamelCase__    )
        create_all_state(i + 1    ,  UpperCamelCase__    ,  level - 1    ,  UpperCamelCase__    ,  UpperCamelCase__    )
        current_list.pop()






def __lowerCAmelCase			(							UpperCamelCase__    )     ->   None:
    for i in total_list:
        print(*UpperCamelCase__    )


if __name__ == "__main__":
  __UpperCAmelCase												=4
  __UpperCAmelCase												=2
  __UpperCAmelCase												=generate_all_combinations(n, k)
  print_all_state(total_list)
