





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


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






class    UpperCAmelCase	(       UpperCamelCase__				):
 def __init__(				self      :Optional[Any]						,       lowercase_      :Dict						,       lowercase_      :List[Any]						,       lowercase_      :Optional[int]=None						,       lowercase_      :List[str]=1							)->							int:
     A__				=			tokenizer
     A__				=			dataset
     A__				=			len(lowercase_							) if n_tasks is None else n_tasks
     A__				=			n_copies



 def __iter__(				self      :Tuple							)->							str:
     A__				=			[]
     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()							)
     A__				=			self.tokenizer(lowercase_						,       padding=lowercase_						,       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	(       UpperCamelCase__				):
 def __init__(				self      :List[Any]						,       lowercase_      :str						,       lowercase_      :Optional[Any]						,       lowercase_      :Optional[Any]							)->							int:
     A__				=			start_length
     A__				=			eof_strings
     A__				=			tokenizer



 def __call__(				self      :str						,       lowercase_      :List[Any]						,       lowercase_      :Dict						,       **lowercase_      :str							)->							int:
     A__				=			self.tokenizer.batch_decode(input_ids[:, self.start_length :]							)
     A__				=			[]
     for decoded_generation in decoded_generations:
         done.append(any(stop_string in decoded_generation for stop_string in self.eof_strings							)							)
     return all(lowercase_							)





def 						UpperCamelCase				(							_lowerCamelCase					:					Optional[int]       ):
    A__				=			re.split("(%s)" % "|".join(_lowerCamelCase       )						,       _lowerCamelCase       )
    # last string should be ""
    return "".join(string_list[:-2]       )





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

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

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

    A__				=			[[] for _ in range(_lowerCamelCase       )]
    for task, generated_tokens in gen_token_dict.items():
        for s in generated_tokens:
            A__				=			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				(							):
    # Setup configuration
    A__				=			HfArgumentParser(_lowerCamelCase       )
    A__				=			parser.parse_args()

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

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

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

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

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

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

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

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

    A__,	A__				=			accelerator.prepare(_lowerCamelCase						,       _lowerCamelCase       )

    A__				=			complete_code(
        _lowerCamelCase						,       _lowerCamelCase						,       _lowerCamelCase						,       _lowerCamelCase						,       n_tasks=_lowerCamelCase						,       batch_size=args.batch_size						,       **_lowerCamelCase						,       )

    if accelerator.is_main_process:
        A__				=			[]

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

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










'''simple docstring'''

import torch

from diffusers import DPMSolverSDEScheduler
from diffusers.utils import torch_device
from diffusers.utils.testing_utils import require_torchsde

from .test_schedulers import SchedulerCommonTest






@require_torchsde
class    UpperCAmelCase	(       UpperCamelCase__				):
 __lowercase   		=				(DPMSolverSDEScheduler,)
 __lowercase   		=				10
 def      UpperCAmelCase_							(				self      :List[Any]						,       **lowercase_      :Optional[int]							)->							str:
     A__				=			{
         "num_train_timesteps": 11_00,
         "beta_start": 0.0_0_0_1,
         "beta_end": 0.0_2,
         "beta_schedule": "linear",
         "noise_sampler_seed": 0,
     }

     config.update(**lowercase_							)
     return config
 def      UpperCAmelCase_							(				self      :int							)->							Dict:
     for timesteps in [10, 50, 1_00, 10_00]:
         self.check_over_configs(num_train_timesteps=lowercase_							)
 def      UpperCAmelCase_							(				self      :List[Any]							)->							Tuple:
     for beta_start, beta_end in zip([0.0_0_0_0_1, 0.0_0_0_1, 0.0_0_1]						,       [0.0_0_0_2, 0.0_0_2, 0.0_2]							):
         self.check_over_configs(beta_start=lowercase_						,       beta_end=lowercase_							)
 def      UpperCAmelCase_							(				self      :Any							)->							Optional[Any]:
     for schedule in ["linear", "scaled_linear"]:
         self.check_over_configs(beta_schedule=lowercase_							)
 def      UpperCAmelCase_							(				self      :List[Any]							)->							Dict:
     for prediction_type in ["epsilon", "v_prediction"]:
         self.check_over_configs(prediction_type=lowercase_							)
 def      UpperCAmelCase_							(				self      :List[str]							)->							Union[str, Any]:
     A__				=			self.scheduler_classes[0]
     A__				=			self.get_scheduler_config()
     A__				=			scheduler_class(**lowercase_							)

     scheduler.set_timesteps(self.num_inference_steps							)

     A__				=			self.dummy_model()
     A__				=			self.dummy_sample_deter * scheduler.init_noise_sigma
     A__				=			sample.to(lowercase_							)

     for i, t in enumerate(scheduler.timesteps							):
         A__				=			scheduler.scale_model_input(lowercase_						,       lowercase_							)

         A__				=			model(lowercase_						,       lowercase_							)

         A__				=			scheduler.step(lowercase_						,       lowercase_						,       lowercase_							)
         A__				=			output.prev_sample

     A__				=			torch.sum(torch.abs(lowercase_							)							)
     A__				=			torch.mean(torch.abs(lowercase_							)							)

     if torch_device in ["mps"]:
         assert abs(result_sum.item() - 1_6_7.4_7_8_2_1_0_4_4_9_2_1_8_7_5							) < 1E-2
         assert abs(result_mean.item() - 0.2_1_7_8_7_0_5_9_6_4_5_6_5_2_7_7							) < 1E-3
     elif torch_device in ["cuda"]:
         assert abs(result_sum.item() - 1_7_1.5_9_3_5_2_1_1_1_8_1_6_4_0_6							) < 1E-2
         assert abs(result_mean.item() - 0.2_2_3_4_2_9_0_6_8_9_2_2_9_9_6_5_2							) < 1E-3
     else:
         assert abs(result_sum.item() - 1_6_2.5_2_3_8_3_4_2_2_8_5_1_5_6_2							) < 1E-2
         assert abs(result_mean.item() - 0.2_1_1_6_1_9_5_7_0_8_5_1_3_2_6							) < 1E-3
 def      UpperCAmelCase_							(				self      :Optional[int]							)->							Dict:
     A__				=			self.scheduler_classes[0]
     A__				=			self.get_scheduler_config(prediction_type="v_prediction"							)
     A__				=			scheduler_class(**lowercase_							)

     scheduler.set_timesteps(self.num_inference_steps							)

     A__				=			self.dummy_model()
     A__				=			self.dummy_sample_deter * scheduler.init_noise_sigma
     A__				=			sample.to(lowercase_							)

     for i, t in enumerate(scheduler.timesteps							):
         A__				=			scheduler.scale_model_input(lowercase_						,       lowercase_							)

         A__				=			model(lowercase_						,       lowercase_							)

         A__				=			scheduler.step(lowercase_						,       lowercase_						,       lowercase_							)
         A__				=			output.prev_sample

     A__				=			torch.sum(torch.abs(lowercase_							)							)
     A__				=			torch.mean(torch.abs(lowercase_							)							)

     if torch_device in ["mps"]:
         assert abs(result_sum.item() - 1_2_4.7_7_1_4_9_2_0_0_4_3_9_4_5_3							) < 1E-2
         assert abs(result_mean.item() - 0.1_6_2_2_6_2_8_9_0_1_4_8_1_6_2_8_4							) < 1E-3
     elif torch_device in ["cuda"]:
         assert abs(result_sum.item() - 1_2_8.1_6_6_3_3_6_0_5_9_5_7_0_3							) < 1E-2
         assert abs(result_mean.item() - 0.1_6_6_8_8_3_2_6_0_0_1_1_6_7_2_9_7							) < 1E-3
     else:
         assert abs(result_sum.item() - 1_1_9.8_4_8_7_5_4_8_8_2_8_1_2_5							) < 1E-2
         assert abs(result_mean.item() - 0.1_5_6_0_5_3_0_6_6_2_5_3_6_6_2_1							) < 1E-3
 def      UpperCAmelCase_							(				self      :Optional[int]							)->							List[str]:
     A__				=			self.scheduler_classes[0]
     A__				=			self.get_scheduler_config()
     A__				=			scheduler_class(**lowercase_							)

     scheduler.set_timesteps(self.num_inference_steps						,       device=lowercase_							)

     A__				=			self.dummy_model()
     A__				=			self.dummy_sample_deter.to(lowercase_							) * scheduler.init_noise_sigma

     for t in scheduler.timesteps:
         A__				=			scheduler.scale_model_input(lowercase_						,       lowercase_							)

         A__				=			model(lowercase_						,       lowercase_							)

         A__				=			scheduler.step(lowercase_						,       lowercase_						,       lowercase_							)
         A__				=			output.prev_sample

     A__				=			torch.sum(torch.abs(lowercase_							)							)
     A__				=			torch.mean(torch.abs(lowercase_							)							)

     if torch_device in ["mps"]:
         assert abs(result_sum.item() - 1_6_7.4_6_9_5_7_3_9_7_4_6_0_9_3_8							) < 1E-2
         assert abs(result_mean.item() - 0.2_1_8_0_5_9_3_4_6_0_7_9_8_2_6_3_5							) < 1E-3
     elif torch_device in ["cuda"]:
         assert abs(result_sum.item() - 1_7_1.5_9_3_5_3_6_3_7_6_9_5_3_1_2							) < 1E-2
         assert abs(result_mean.item() - 0.2_2_3_4_2_9_0_8_3_8_2_4_1_5_7_7_1							) < 1E-3
     else:
         assert abs(result_sum.item() - 1_6_2.5_2_3_8_3_4_2_2_8_5_1_5_6_2							) < 1E-2
         assert abs(result_mean.item() - 0.2_1_1_6_1_9_5_7_0_8_5_1_3_2_6							) < 1E-3



 def      UpperCAmelCase_							(				self      :Tuple							)->							Dict:
     A__				=			self.scheduler_classes[0]
     A__				=			self.get_scheduler_config()
     A__				=			scheduler_class(**lowercase_						,       use_karras_sigmas=lowercase_							)

     scheduler.set_timesteps(self.num_inference_steps						,       device=lowercase_							)

     A__				=			self.dummy_model()
     A__				=			self.dummy_sample_deter.to(lowercase_							) * scheduler.init_noise_sigma
     A__				=			sample.to(lowercase_							)

     for t in scheduler.timesteps:
         A__				=			scheduler.scale_model_input(lowercase_						,       lowercase_							)

         A__				=			model(lowercase_						,       lowercase_							)

         A__				=			scheduler.step(lowercase_						,       lowercase_						,       lowercase_							)
         A__				=			output.prev_sample

     A__				=			torch.sum(torch.abs(lowercase_							)							)
     A__				=			torch.mean(torch.abs(lowercase_							)							)

     if torch_device in ["mps"]:
         assert abs(result_sum.item() - 1_7_6.6_6_9_7_4_1_3_5_7_4_2_1_8_8							) < 1E-2
         assert abs(result_mean.item() - 0.2_3_0_0_3_8_7_2_7_3_0_9_8_1_8_1_1							) < 1E-2
     elif torch_device in ["cuda"]:
         assert abs(result_sum.item() - 1_7_7.6_3_6_5_3_5_6_4_4_5_3_1_2_5							) < 1E-2
         assert abs(result_mean.item() - 0.2_3_0_0_3_8_7_2_7_3_0_9_8_1_8_1_1							) < 1E-2
     else:
         assert abs(result_sum.item() - 1_7_0.3_1_3_5_2_2_3_3_8_8_6_7_2							) < 1E-2
         assert abs(result_mean.item() - 0.2_3_0_0_3_8_7_2_7_3_0_9_8_1_8_1_1							) < 1E-2



