
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


__magic_name__            =  ['''\nclass''', '''\ndef''', '''\n#''', '''\n@''', '''\nprint''', '''\nif''']




class 					_SCREAMING_SNAKE_CASE  (      __UpperCamelCase ):
      def __init__(			self		,	lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase=None		,	lowerCamelCase=1 ):
          snake_case__						=				tokenizer
          snake_case__						=				dataset
          snake_case__						=				len(lowerCamelCase ) if n_tasks is None else n_tasks
          snake_case__						=				n_copies



      def __iter__(			self ):
          snake_case__						=				[]
          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() )
          snake_case__						=				self.tokenizer(lowerCamelCase		,	padding=lowerCamelCase		,	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 					_SCREAMING_SNAKE_CASE  (      __UpperCamelCase ):
      def __init__(			self		,	lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase ):
          snake_case__						=				start_length
          snake_case__						=				eof_strings
          snake_case__						=				tokenizer



      def __call__(			self		,	lowerCamelCase		,	lowerCamelCase		,	**lowerCamelCase ):
          snake_case__						=				self.tokenizer.batch_decode(input_ids[:, self.start_length :] )
          snake_case__						=				[]
          for decoded_generation in decoded_generations:
              done.append(any(stop_string in decoded_generation for stop_string in self.eof_strings ) )
          return all(lowerCamelCase )





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





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

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

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

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





def 	SCREAMING_SNAKE_CASE__    (      ):
    # Setup configuration
    snake_case__						=				HfArgumentParser(__lowerCAmelCase							)
    snake_case__						=				parser.parse_args()

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

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

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

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

    # Generation settings
    snake_case__						=				{
        "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
    snake_case__						=				load_dataset("openai_humaneval"							)
    snake_case__						=				load_metric("code_eval"							)

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

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

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

    snake_case__						,					snake_case__						=				accelerator.prepare(__lowerCAmelCase    ,		__lowerCAmelCase							)

    snake_case__						=				complete_code(
        __lowerCAmelCase    ,		__lowerCAmelCase    ,		__lowerCAmelCase    ,		__lowerCAmelCase    ,		n_tasks=__lowerCAmelCase    ,		batch_size=args.batch_size    ,		**__lowerCAmelCase    ,		)

    if accelerator.is_main_process:
        snake_case__						=				[]

        for task in tqdm(range(__lowerCAmelCase							)							):
            snake_case__						=				human_eval["test"][task]["test"]
            snake_case__						=				F"""check({human_eval["test"][task]["entry_point"]})"""
            references.append("\n" + test_func + "\n" + entry_point							)

        # Evaluate completions with "code_eval" metric
        snake_case__						,					snake_case__						=				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()





def 	SCREAMING_SNAKE_CASE__    (      __lowerCAmelCase    ,		__lowerCAmelCase							):
    return price * (1 + tax_rate)


if __name__ == "__main__":
     print(F'''{price_plus_tax(100, 0.2_5) = }''')
     print(F'''{price_plus_tax(1_2_5.5_0, 0.0_5) = }''')



