





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


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




class     SCREAMING_SNAKE_CASE__ (						UpperCamelCase__      ):



 def __init__(							self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase=None,__lowerCamelCase=1 ):
        A__											=						tokenizer
        A__											=						dataset
        A__											=						len(__lowerCamelCase ) if n_tasks is None else n_tasks
        A__											=						n_copies




 def __iter__(							self ):
        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(__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 ):
        A__											=						start_length
        A__											=						eof_strings
        A__											=						tokenizer




 def __call__(							self,__lowerCamelCase,__lowerCamelCase,**__lowerCamelCase ):
        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(__lowerCamelCase )


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


def 	UpperCamelCase__(       UpperCamelCase__			:       Optional[Any] ,     UpperCamelCase__			:       str ,     UpperCamelCase__			:       Union[str, Any] ,     UpperCamelCase__			:       List[str] ,     UpperCamelCase__			:       int ,     UpperCamelCase__			:       Tuple=20 ,     **UpperCamelCase__			:       Tuple  )->str:
       A__											=						defaultdict(UpperCamelCase__  )  # dict of list of generated tokens
       for step, batch in tqdm(enumerate(UpperCamelCase__  )  ):
              with torch.no_grad():
                     A__											=						batch['''ids'''].shape[-1]
                     A__											=						accelerator.unwrap_model(UpperCamelCase__  ).generate(
                         input_ids=batch['''ids'''][:, : batch['''input_len''']] ,     num_return_sequences=UpperCamelCase__ ,     **UpperCamelCase__  )
                     # each task is generated batch_size times
                     A__											=						batch['''task_id'''].repeat(UpperCamelCase__  )
                     A__											=						accelerator.pad_across_processes(
                         UpperCamelCase__ ,     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(UpperCamelCase__ ,     UpperCamelCase__  ):
                            gen_token_dict[task].append(UpperCamelCase__  )

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


def 	UpperCamelCase__(       )->List[str]:
       # Setup configuration
       A__											=						HfArgumentParser(UpperCamelCase__  )
       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=UpperCamelCase__  )

       # 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 ,     UpperCamelCase__ ,     UpperCamelCase__  )]  ),
       }

       # 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(UpperCamelCase__ ,     human_eval['''test'''] ,     n_copies=UpperCamelCase__ ,     n_tasks=UpperCamelCase__  )
       # do not confuse args.batch_size, which is actually the num_return_sequences
       A__											=						DataLoader(UpperCamelCase__ ,     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(UpperCamelCase__ ,     UpperCamelCase__  )

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

       if accelerator.is_main_process:
              A__											=						[]

              for task in tqdm(range(UpperCamelCase__  )  ):
                     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=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()












import unittest

from huggingface_hub import hf_hub_download

from transformers import MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING, VideoMAEFeatureExtractor
from transformers.pipelines import VideoClassificationPipeline, pipeline
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_decord,
    require_tf,
    require_torch,
    require_torch_or_tf,
    require_vision,
)

from .test_pipelines_common import ANY




@is_pipeline_test
@require_torch_or_tf
@require_vision
@require_decord
class     SCREAMING_SNAKE_CASE__ (						unittest.TestCase      ):
 __SCREAMING_SNAKE_CASE							       =	MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING



 def 				UpperCamelCase			(							self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase ):
        A__											=						hf_hub_download(
            repo_id='''nateraw/video-demo''',filename='''archery.mp4''',repo_type='''dataset''' )
        A__											=						VideoClassificationPipeline(model=__lowerCamelCase,image_processor=__lowerCamelCase,top_k=2 )
        A__											=						[
            example_video_filepath,
            '''https://huggingface.co/datasets/nateraw/video-demo/resolve/main/archery.mp4''',
        ]
        return video_classifier, examples



 def 				UpperCamelCase			(							self,__lowerCamelCase,__lowerCamelCase ):
        for example in examples:
               A__											=						video_classifier(__lowerCamelCase )

               self.assertEqual(
                   __lowerCamelCase,[
                       {'''score''': ANY(__lowerCamelCase ), '''label''': ANY(__lowerCamelCase )},
                       {'''score''': ANY(__lowerCamelCase ), '''label''': ANY(__lowerCamelCase )},
                   ],)



 @require_torch
 def 				UpperCamelCase			(							self ):
        A__											=						'''hf-internal-testing/tiny-random-VideoMAEForVideoClassification'''
        A__											=						VideoMAEFeatureExtractor(
            size={'''shortest_edge''': 10},crop_size={'''height''': 10, '''width''': 10} )
        A__											=						pipeline(
            '''video-classification''',model=__lowerCamelCase,feature_extractor=__lowerCamelCase,frame_sampling_rate=4 )

        A__											=						hf_hub_download(repo_id='''nateraw/video-demo''',filename='''archery.mp4''',repo_type='''dataset''' )
        A__											=						video_classifier(__lowerCamelCase,top_k=2 )
        self.assertEqual(
            nested_simplify(__lowerCamelCase,decimals=4 ),[{'''score''': 0.5199, '''label''': '''LABEL_0'''}, {'''score''': 0.4801, '''label''': '''LABEL_1'''}],)

        A__											=						video_classifier(
            [
                video_file_path,
                video_file_path,
            ],top_k=2,)
        self.assertEqual(
            nested_simplify(__lowerCamelCase,decimals=4 ),[
                [{'''score''': 0.5199, '''label''': '''LABEL_0'''}, {'''score''': 0.4801, '''label''': '''LABEL_1'''}],
                [{'''score''': 0.5199, '''label''': '''LABEL_0'''}, {'''score''': 0.4801, '''label''': '''LABEL_1'''}],
            ],)




 @require_tf
 def 				UpperCamelCase			(							self ):
        pass





