



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     lowerCAmelCase_		(  _a							):
       def __init__(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_=None,      SCREAMING_SNAKE_CASE_=1	)					->		List[Any]:
          UpperCamelCase   :  int						     =     tokenizer
          UpperCamelCase   :  Dict						     =     dataset
          UpperCamelCase   :  Optional[int]						     =     len(snake_case_	) if n_tasks is None else n_tasks
          UpperCamelCase   :  List[str]						     =     n_copies






       def __iter__(      self	)					->		Tuple:
          UpperCamelCase   :  Optional[int]						     =     []
          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()	)
          UpperCamelCase   :  Optional[int]						     =     self.tokenizer(snake_case_,      padding=snake_case_,      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     lowerCAmelCase_		(  _a							):
       def __init__(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	)					->		Optional[Any]:
          UpperCamelCase   :  Dict						     =     start_length
          UpperCamelCase   :  Optional[Any]						     =     eof_strings
          UpperCamelCase   :  List[str]						     =     tokenizer






       def __call__(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      **SCREAMING_SNAKE_CASE_	)					->		str:
          UpperCamelCase   :  int						     =     self.tokenizer.batch_decode(input_ids[:, self.start_length :]	)
          UpperCamelCase   :  Optional[int]						     =     []
          for decoded_generation in decoded_generations:
             done.append(any(stop_string in decoded_generation for stop_string in self.eof_strings	)	)
          return all(snake_case_	)






def 		UpperCamelCase     (       snake_case__	:   Optional[Any]     )     ->					Union[str, Any]:
   UpperCamelCase   :  List[Any]						     =     re.split('(%s)' % '|'.join(_lowerCAmelCase     )	,					_lowerCAmelCase     )
   # last string should be ""
   return "".join(string_list[:-2]     )
def 		UpperCamelCase     (       snake_case__	:   Optional[int]	,					snake_case__	:   Union[str, Any]	,					snake_case__	:   Optional[Any]	,					snake_case__	:   List[Any]	,					snake_case__	:   Optional[int]	,					snake_case__	:   Optional[Any]=20	,					**snake_case__	:   str     )     ->					Union[str, Any]:
   UpperCamelCase   :  Optional[Any]						     =     defaultdict(_lowerCAmelCase     )  # dict of list of generated tokens
   for step, batch in tqdm(enumerate(_lowerCAmelCase     )     ):
      with torch.no_grad():
         UpperCamelCase   :  Optional[int]						     =     batch["""ids"""].shape[-1]
         UpperCamelCase   :  List[Any]						     =     accelerator.unwrap_model(_lowerCAmelCase     ).generate(
             input_ids=batch['ids'][:, : batch['input_len']]	,					num_return_sequences=_lowerCAmelCase	,					**_lowerCAmelCase     )
         # each task is generated batch_size times
         UpperCamelCase   :  Optional[int]						     =     batch["""task_id"""].repeat(_lowerCAmelCase     )
         UpperCamelCase   :  List[Any]						     =     accelerator.pad_across_processes(
             _lowerCAmelCase	,					dim=1	,					pad_index=tokenizer.pad_token_id     )

         UpperCamelCase   :  Optional[Any]						     =     accelerator.gather((generated_tokens, generated_tasks)     )
         UpperCamelCase   :  Optional[Any]						     =     generated_tokens.cpu().numpy()
         UpperCamelCase   :  Optional[int]						     =     generated_tasks.cpu().numpy()

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

   UpperCamelCase   :  Optional[Any]						     =     [[] for _ in range(_lowerCAmelCase     )]
   for task, generated_tokens in gen_token_dict.items():
      for s in generated_tokens:
         UpperCamelCase   :  int						     =     tokenizer.decode(_lowerCAmelCase	,					skip_special_tokens=_lowerCAmelCase	,					clean_up_tokenization_spaces=_lowerCAmelCase     )
         code_gens[task].append(remove_last_block(_lowerCAmelCase     )     )
   return code_gens
def 		UpperCamelCase     (       )     ->					Union[str, Any]:
   # Setup configuration
   UpperCamelCase   :  List[str]						     =     HfArgumentParser(_lowerCAmelCase     )
   UpperCamelCase   :  Optional[int]						     =     parser.parse_args()

   transformers.logging.set_verbosity_error()
   # enables code execution in code_eval metric
   UpperCamelCase   :  Union[str, Any]						     =     args.HF_ALLOW_CODE_EVAL
   # make sure tokenizer plays nice with multiprocessing
   UpperCamelCase   :  Union[str, Any]						     =     """false"""

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

   # Use dataset load to feed to accelerate
   UpperCamelCase   :  int						     =     Accelerator()
   set_seed(args.seed	,					device_specific=_lowerCAmelCase     )

   # Load model and tokenizer
   UpperCamelCase   :  Dict						     =     AutoTokenizer.from_pretrained(args.model_ckpt     )
   UpperCamelCase   :  Dict						     =     tokenizer.eos_token
   UpperCamelCase   :  Any						     =     AutoModelForCausalLM.from_pretrained(args.model_ckpt     )

   # Generation settings
   UpperCamelCase   :  Any						     =     {
       """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	,					_lowerCAmelCase	,					_lowerCAmelCase     )]     ),
   }

   # Load evaluation dataset and metric
   UpperCamelCase   :  Optional[int]						     =     load_dataset('openai_humaneval'     )
   UpperCamelCase   :  List[Any]						     =     load_metric('code_eval'     )

   UpperCamelCase   :  Optional[Any]						     =     args.num_tasks if args.num_tasks is not None else len(human_eval['test']     )
   UpperCamelCase   :  Dict						     =     args.n_samples // args.batch_size

   UpperCamelCase   :  int						     =     TokenizedDataset(_lowerCAmelCase	,					human_eval['test']	,					n_copies=_lowerCAmelCase	,					n_tasks=_lowerCAmelCase     )
   # do not confuse args.batch_size, which is actually the num_return_sequences
   UpperCamelCase   :  Optional[Any]						     =     DataLoader(_lowerCAmelCase	,					batch_size=1     )

   # Run a quick test to see if code evaluation is enabled
   try:
      UpperCamelCase   :  List[Any]						     =     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

   UpperCamelCase   :  Dict						     =     accelerator.prepare(_lowerCAmelCase	,					_lowerCAmelCase     )

   UpperCamelCase   :  str						     =     complete_code(
       _lowerCAmelCase	,					_lowerCAmelCase	,					_lowerCAmelCase	,					_lowerCAmelCase	,					n_tasks=_lowerCAmelCase	,					batch_size=args.batch_size	,					**_lowerCAmelCase	,					)

   if accelerator.is_main_process:
      UpperCamelCase   :  str						     =     []

      for task in tqdm(range(_lowerCAmelCase     )     ):
         UpperCamelCase   :  Any						     =     human_eval["""test"""][task]["""test"""]
         UpperCamelCase   :  int						     =     F"""check({human_eval["test"][task]["entry_point"]})"""
         references.append('\n' + test_func + '\n' + entry_point     )

      # Evaluate completions with "code_eval" metric
      UpperCamelCase   :  Optional[Any]						     =     code_eval_metric.compute(
          references=_lowerCAmelCase	,					predictions=_lowerCAmelCase	,					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(_lowerCAmelCase	,					_lowerCAmelCase     )


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








import argparse
import shlex

import runhouse as rh


if __name__ == "__main__":
    # Refer to https://runhouse-docs.readthedocs-hosted.com/en/latest/api/python/cluster.html#hardware-setup for cloud access
    # setup instructions, if using on-demand hardware

    # If user passes --user <user> --host <host> --key_path <key_path> <example> <args>, fill them in as BYO cluster
    # If user passes --instance <instance> --provider <provider> <example> <args>, fill them in as on-demand cluster
    # Throw an error if user passes both BYO and on-demand cluster args
    # Otherwise, use default values
    __UpperCAmelCase	       =	argparse.ArgumentParser()
    parser.add_argument('''--user''', type=str, default='''ubuntu''')
    parser.add_argument('''--host''', type=str, default='''localhost''')
    parser.add_argument('''--key_path''', type=str, default=None)
    parser.add_argument('''--instance''', type=str, default='''V100:1''')
    parser.add_argument('''--provider''', type=str, default='''cheapest''')
    parser.add_argument('''--use_spot''', type=bool, default=False)
    parser.add_argument('''--example''', type=str, default='''pytorch/text-generation/run_generation.py''')
    __UpperCAmelCase       ,  __UpperCAmelCase	       =	parser.parse_known_args()
    if args.host != "localhost":
        if args.instance != "V100:1" or args.provider != "cheapest":
            raise ValueError('''Cannot specify both BYO and on-demand cluster args''')
        __UpperCAmelCase	       =	rh.cluster(
            name='''rh-cluster''', ips=[args.host], ssh_creds={'''ssh_user''': args.user, '''ssh_private_key''': args.key_path}
        )
    else:
        __UpperCAmelCase	       =	rh.cluster(
            name='''rh-cluster''', instance_type=args.instance, provider=args.provider, use_spot=args.use_spot
        )
    __UpperCAmelCase	       =	args.example.rsplit('''/''', 1)[0]

    # Set up remote environment
    cluster.install_packages(['''pip:./'''])  # Installs transformers from local source
    # Note transformers is copied into the home directory on the remote machine, so we can install from there
    cluster.run([F"""pip install -r transformers/examples/{example_dir}/requirements.txt"""])
    cluster.run(['''pip install torch --upgrade --extra-index-url https://download.pytorch.org/whl/cu117'''])

    # Run example. You can bypass the CLI wrapper and paste your own code here.
    cluster.run([F"""python transformers/examples/{args.example} {" ".join(shlex.quote(arg) for arg in unknown)}"""])

    # Alternatively, we can just import and run a training function (especially if there's no wrapper CLI):
    # from my_script... import train
    # reqs = ['pip:./', 'torch', 'datasets', 'accelerate', 'evaluate', 'tqdm', 'scipy', 'scikit-learn', 'tensorboard']
    # launch_train_gpu = rh.function(fn=train,
    #                                system=gpu,
    #                                reqs=reqs,
    #                                name='train_bert_glue')
    #
    # We can pass in arguments just like we would to a function:
    # launch_train_gpu(num_epochs = 3, lr = 2e-5, seed = 42, batch_size = 16
    #                  stream_logs=True)



