




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


SCREAMING_SNAKE_CASE__       :      List[Any]  =   ['\nclass', '\ndef', '\n#', '\n@', '\nprint', '\nif']




class 		UpperCamelCase__	(lowerCAmelCase__     ):

    '''simple docstring'''



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







    def __iter__(    self     )     ->      Tuple:
           lowerCamelCase   :      Tuple			 =					[]
           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   :      List[str]			 =					self.tokenizer(UpperCamelCase__     ,  padding=UpperCamelCase__     ,  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__	(lowerCAmelCase__     ):

    '''simple docstring'''



    def __init__(    self     ,  UpperCamelCase__     ,  UpperCamelCase__     ,  UpperCamelCase__     )     ->      int:
           lowerCamelCase   :      str			 =					start_length
           lowerCamelCase   :      Tuple			 =					eof_strings
           lowerCamelCase   :      List[Any]			 =					tokenizer







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




def      A      (      _SCREAMING_SNAKE_CASE    )	-> Any:
       lowerCamelCase   :      List[Any]			 =					re.split("(%s)" % "|".join(_SCREAMING_SNAKE_CASE    ) ,_SCREAMING_SNAKE_CASE    )
       # last string should be ""
       return "".join(string_list[:-2]    )

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

                     lowerCamelCase ,      lowerCamelCase   :      Any			 =					accelerator.gather((generated_tokens, generated_tasks)    )
                     lowerCamelCase   :      List[Any]			 =					generated_tokens.cpu().numpy()
                     lowerCamelCase   :      Tuple			 =					generated_tasks.cpu().numpy()

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

       lowerCamelCase   :      str			 =					[[] for _ in range(_SCREAMING_SNAKE_CASE    )]
       for task, generated_tokens in gen_token_dict.items():
              for s in generated_tokens:
                     lowerCamelCase   :      List[Any]			 =					tokenizer.decode(_SCREAMING_SNAKE_CASE ,skip_special_tokens=_SCREAMING_SNAKE_CASE ,clean_up_tokenization_spaces=_SCREAMING_SNAKE_CASE    )
                     code_gens[task].append(remove_last_block(_SCREAMING_SNAKE_CASE    )    )
       return code_gens

def      A      (      )	-> Tuple:
       # Setup configuration
       lowerCamelCase   :      str			 =					HfArgumentParser(_SCREAMING_SNAKE_CASE    )
       lowerCamelCase   :      Union[str, Any]			 =					parser.parse_args()

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

       if args.num_workers is None:
              lowerCamelCase   :      Optional[int]			 =					multiprocessing.cpu_count()

       # Use dataset load to feed to accelerate
       lowerCamelCase   :      Dict			 =					Accelerator()
       set_seed(args.seed ,device_specific=_SCREAMING_SNAKE_CASE    )

       # Load model and tokenizer
       lowerCamelCase   :      Tuple			 =					AutoTokenizer.from_pretrained(args.model_ckpt    )
       lowerCamelCase   :      str			 =					tokenizer.eos_token
       lowerCamelCase   :      Dict			 =					AutoModelForCausalLM.from_pretrained(args.model_ckpt    )

       # Generation settings
       lowerCamelCase   :      List[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 ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE    )]    ),
       }

       # Load evaluation dataset and metric
       lowerCamelCase   :      Optional[Any]			 =					load_dataset("openai_humaneval"    )
       lowerCamelCase   :      str			 =					load_metric("code_eval"    )

       lowerCamelCase   :      str			 =					args.num_tasks if args.num_tasks is not None else len(human_eval["test"]    )
       lowerCamelCase   :      List[Any]			 =					args.n_samples // args.batch_size

       lowerCamelCase   :      Union[str, Any]			 =					TokenizedDataset(_SCREAMING_SNAKE_CASE ,human_eval["test"] ,n_copies=_SCREAMING_SNAKE_CASE ,n_tasks=_SCREAMING_SNAKE_CASE    )
       # do not confuse args.batch_size, which is actually the num_return_sequences
       lowerCamelCase   :      str			 =					DataLoader(_SCREAMING_SNAKE_CASE ,batch_size=1    )

       # Run a quick test to see if code evaluation is enabled
       try:
              lowerCamelCase   :      Dict			 =					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   :      int			 =					accelerator.prepare(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE    )

       lowerCamelCase   :      List[Any]			 =					complete_code(
           _SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,n_tasks=_SCREAMING_SNAKE_CASE ,batch_size=args.batch_size ,**_SCREAMING_SNAKE_CASE ,)

       if accelerator.is_main_process:
              lowerCamelCase   :      List[str]			 =					[]

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

              # Evaluate completions with "code_eval" metric
              lowerCamelCase ,      lowerCamelCase   :      Dict			 =					code_eval_metric.compute(
                  references=_SCREAMING_SNAKE_CASE ,predictions=_SCREAMING_SNAKE_CASE ,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(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE    )


# 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
from pathlib import Path

from transformers import AutoConfig, AutoTokenizer, RagConfig, RagSequenceForGeneration, RagTokenForGeneration

def      A      (      _SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE = None ,_SCREAMING_SNAKE_CASE = None ,_SCREAMING_SNAKE_CASE = None ,)	-> List[str]:
       if config_name_or_path is None:
              lowerCamelCase   :      Any			 =					"facebook/rag-token-base" if model_type == "rag_token" else "facebook/rag-sequence-base"

       if generator_tokenizer_name_or_path is None:
              lowerCamelCase   :      Dict			 =					generator_name_or_path

       if question_encoder_tokenizer_name_or_path is None:
              lowerCamelCase   :      Any			 =					question_encoder_name_or_path

       lowerCamelCase   :      str			 =					RagTokenForGeneration if model_type == "rag_token" else RagSequenceForGeneration

       # Save model.
       lowerCamelCase   :      List[Any]			 =					RagConfig.from_pretrained(_SCREAMING_SNAKE_CASE    )
       lowerCamelCase   :      Union[str, Any]			 =					AutoConfig.from_pretrained(_SCREAMING_SNAKE_CASE    )
       lowerCamelCase   :      Optional[int]			 =					AutoConfig.from_pretrained(_SCREAMING_SNAKE_CASE    )

       lowerCamelCase   :      Optional[Any]			 =					gen_config
       lowerCamelCase   :      Optional[Any]			 =					question_encoder_config

       lowerCamelCase   :      List[Any]			 =					model_class.from_pretrained_question_encoder_generator(
           _SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,config=_SCREAMING_SNAKE_CASE    )
       rag_model.save_pretrained(_SCREAMING_SNAKE_CASE    )

       # Sanity check.
       model_class.from_pretrained(_SCREAMING_SNAKE_CASE    )

       # Save tokenizers.
       lowerCamelCase   :      List[str]			 =					AutoTokenizer.from_pretrained(_SCREAMING_SNAKE_CASE    )
       gen_tokenizer.save_pretrained(dest_dir / "generator_tokenizer/"    )
       lowerCamelCase   :      int			 =					AutoTokenizer.from_pretrained(_SCREAMING_SNAKE_CASE    )
       question_encoder_tokenizer.save_pretrained(dest_dir / "question_encoder_tokenizer/"    )


if __name__ == "__main__":
       SCREAMING_SNAKE_CASE__       :      Any  =   argparse.ArgumentParser()
       parser.add_argument(
           '--model_type',
           choices=['rag_sequence', 'rag_token'],
           required=True,
           type=str,
           help='RAG model type: rag_sequence, rag_token',
       )
       parser.add_argument('--dest', type=str, required=True, help='Path to the output checkpoint directory.')
       parser.add_argument('--generator_name_or_path', type=str, required=True, help='Generator model identifier')
       parser.add_argument(
           '--question_encoder_name_or_path', type=str, required=True, help='Question encoder model identifier'
       )

       parser.add_argument(
           '--generator_tokenizer_name_or_path',
           type=str,
           help='Generator tokenizer identifier, if not specified, resolves to ``generator_name_or_path``',
       )
       parser.add_argument(
           '--question_encoder_tokenizer_name_or_path',
           type=str,
           help='Question encoder tokenizer identifier, if not specified, resolves to ``question_encoder_name_or_path``',
       )
       parser.add_argument(
           '--config_name_or_path',
           type=str,
           help=(
               'Identifier of the model config to use, if not provided, resolves to a base config for a given'
               ' ``model_type``'
           ),
       )

       SCREAMING_SNAKE_CASE__       :      Union[str, Any]  =   parser.parse_args()

       SCREAMING_SNAKE_CASE__       :      Optional[Any]  =   Path(args.dest)
       dest_dir.mkdir(exist_ok=True)

       consolidate(
           args.model_type,
           args.generator_name_or_path,
           args.question_encoder_name_or_path,
           dest_dir,
           args.config_name_or_path,
           args.generator_tokenizer_name_or_path,
           args.question_encoder_tokenizer_name_or_path,
       )






