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




import os
import shutil
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
from datasets import Dataset

from transformers.models.realm.configuration_realm import RealmConfig
from transformers.models.realm.retrieval_realm import _REALM_BLOCK_RECORDS_FILENAME, RealmRetriever
from transformers.models.realm.tokenization_realm import VOCAB_FILES_NAMES, RealmTokenizer




class    a__				(       UpperCAmelCase__			):






     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[Any]   ):





         """simple docstring"""



         __lowerCamelCase          =  tempfile.mkdtemp()
         __lowerCamelCase          =  5

         # Realm tok
         __lowerCamelCase          =  [
             '''[UNK]''',
             '''[CLS]''',
             '''[SEP]''',
             '''[PAD]''',
             '''[MASK]''',
             '''test''',
             '''question''',
             '''this''',
             '''is''',
             '''the''',
             '''first''',
             '''second''',
             '''third''',
             '''fourth''',
             '''fifth''',
             '''record''',
             '''want''',
             '''##want''',
             '''##ed''',
             '''wa''',
             '''un''',
             '''runn''',
             '''##ing''',
             ''',''',
             '''low''',
             '''lowest''',
         ]
         __lowerCamelCase          =  os.path.join(self.tmpdirname  ,    '''realm_tokenizer'''   )
         os.makedirs(a  ,    exist_ok=a   )
         __lowerCamelCase          =  os.path.join(a  ,    VOCAB_FILES_NAMES['''vocab_file''']   )
         with open(self.vocab_file  ,    '''w'''  ,    encoding='''utf-8'''   ) as vocab_writer:
             vocab_writer.write(''''''.join([x + '''\n''' for x in vocab_tokens]   )   )

         __lowerCamelCase          =  os.path.join(self.tmpdirname  ,    '''realm_block_records'''   )
         os.makedirs(a  ,    exist_ok=a   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[Any]   ):





         """simple docstring"""



         return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname  ,    '''realm_tokenizer'''   )   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		Any   ):





         """simple docstring"""



         shutil.rmtree(self.tmpdirname   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[Any]   ):





         """simple docstring"""



         __lowerCamelCase          =  RealmConfig(num_block_records=self.num_block_records   )
         return config






     def        SCREAMING_SNAKE_CASE__       (			self		:		int   ):





         """simple docstring"""



         __lowerCamelCase          =  Dataset.from_dict(
             {
                 '''id''': ['''0''', '''1'''],
                 '''question''': ['''foo''', '''bar'''],
                 '''answers''': [['''Foo''', '''Bar'''], ['''Bar''']],
             }   )
         return dataset






     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[Any]   ):





         """simple docstring"""



         __lowerCamelCase          =  np.array(
             [
                 b'''This is the first record''',
                 b'''This is the second record''',
                 b'''This is the third record''',
                 b'''This is the fourth record''',
                 b'''This is the fifth record''',
                 b'''This is a longer longer longer record''',
             ]  ,    dtype=a  ,    )
         return block_records






     def        SCREAMING_SNAKE_CASE__       (			self		:		str   ):





         """simple docstring"""



         __lowerCamelCase          =  RealmRetriever(
             block_records=self.get_dummy_block_records()  ,    tokenizer=self.get_tokenizer()  ,    )
         return retriever






     def        SCREAMING_SNAKE_CASE__       (			self		:		Any   ):





         """simple docstring"""



         __lowerCamelCase          =  self.get_config()
         __lowerCamelCase          =  self.get_dummy_retriever()
         __lowerCamelCase          =  retriever.tokenizer

         __lowerCamelCase          =  np.array([0, 3]  ,    dtype='''long'''   )
         __lowerCamelCase          =  tokenizer(['''Test question''']   ).input_ids
         __lowerCamelCase          =  tokenizer(
             ['''the fourth''']  ,    add_special_tokens=a  ,    return_token_type_ids=a  ,    return_attention_mask=a  ,    ).input_ids
         __lowerCamelCase          =  config.reader_seq_len

         __lowerCamelCase     ,  __lowerCamelCase     ,  __lowerCamelCase     ,  __lowerCamelCase          =  retriever(
             a  ,    a  ,    answer_ids=a  ,    max_length=a  ,    return_tensors='''np'''   )

         self.assertEqual(len(a   )  ,    2   )
         self.assertEqual(len(a   )  ,    2   )
         self.assertEqual(len(a   )  ,    2   )
         self.assertEqual(concat_inputs.input_ids.shape  ,    (2, 10)   )
         self.assertEqual(concat_inputs.attention_mask.shape  ,    (2, 10)   )
         self.assertEqual(concat_inputs.token_type_ids.shape  ,    (2, 10)   )
         self.assertEqual(concat_inputs.special_tokens_mask.shape  ,    (2, 10)   )
         self.assertEqual(
             tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[0]   )  ,    ['''[CLS]''', '''test''', '''question''', '''[SEP]''', '''this''', '''is''', '''the''', '''first''', '''record''', '''[SEP]''']  ,    )
         self.assertEqual(
             tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[1]   )  ,    ['''[CLS]''', '''test''', '''question''', '''[SEP]''', '''this''', '''is''', '''the''', '''fourth''', '''record''', '''[SEP]''']  ,    )






     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[Any]   ):





         """simple docstring"""



         __lowerCamelCase          =  self.get_config()
         __lowerCamelCase          =  self.get_dummy_retriever()
         __lowerCamelCase          =  retriever.tokenizer

         __lowerCamelCase          =  np.array([0, 3, 5]  ,    dtype='''long'''   )
         __lowerCamelCase          =  tokenizer(['''Test question''']   ).input_ids
         __lowerCamelCase          =  tokenizer(
             ['''the fourth''', '''longer longer''']  ,    add_special_tokens=a  ,    return_token_type_ids=a  ,    return_attention_mask=a  ,    ).input_ids
         __lowerCamelCase          =  config.reader_seq_len

         __lowerCamelCase     ,  __lowerCamelCase     ,  __lowerCamelCase     ,  __lowerCamelCase          =  retriever(
             a  ,    a  ,    answer_ids=a  ,    max_length=a  ,    return_tensors='''np'''   )

         self.assertEqual([False, True, True]  ,    a   )
         self.assertEqual([[-1, -1, -1], [6, -1, -1], [6, 7, 8]]  ,    a   )
         self.assertEqual([[-1, -1, -1], [7, -1, -1], [7, 8, 9]]  ,    a   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		int   ):





         """simple docstring"""



         __lowerCamelCase          =  self.get_dummy_retriever()
         retriever.save_pretrained(os.path.join(self.tmpdirname  ,    '''realm_block_records'''   )   )

         # Test local path
         __lowerCamelCase          =  retriever.from_pretrained(os.path.join(self.tmpdirname  ,    '''realm_block_records'''   )   )
         self.assertEqual(retriever.block_records[0]  ,    b'''This is the first record'''   )

         # Test mocked remote path
         with patch('''transformers.models.realm.retrieval_realm.hf_hub_download'''   ) as mock_hf_hub_download:
             __lowerCamelCase          =  os.path.join(
                 os.path.join(self.tmpdirname  ,    '''realm_block_records'''   )  ,    _REALM_BLOCK_RECORDS_FILENAME   )
             __lowerCamelCase          =  RealmRetriever.from_pretrained('''google/realm-cc-news-pretrained-openqa'''   )

         self.assertEqual(retriever.block_records[0]  ,    b'''This is the first record'''   )
