






"""simple docstring"""





import argparse
import collections
import os
import re

from transformers.utils import direct_transformers_import


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/check_table.py
_SCREAMING_SNAKE_CASE					:				Union[str, Any]					    =						'''src/transformers'''
_SCREAMING_SNAKE_CASE					:				Tuple					    =						'''docs/source/en'''
_SCREAMING_SNAKE_CASE					:				str					    =						'''.'''







def      lowerCamelCase__       (							_lowerCamelCase       :		Dict		,				_lowerCamelCase       :		Any		,				_lowerCamelCase       :		Union[str, Any] )      ->				Any:
  with open(_lowerCamelCase		,				'r'		,				encoding='utf-8'		,				newline='\n' ) as f:
    lowerCamelCase_         =						f.readlines()
  # Find the start prompt.
  lowerCamelCase_         =						0
  while not lines[start_index].startswith(_lowerCamelCase ):
    start_index += 1
  start_index += 1

  lowerCamelCase_         =						start_index
  while not lines[end_index].startswith(_lowerCamelCase ):
    end_index += 1
  end_index -= 1

  while len(lines[start_index] ) <= 1:
    start_index += 1
  while len(lines[end_index] ) <= 1:
    end_index -= 1
  end_index += 1
  return "".join(lines[start_index:end_index] ), start_index, end_index, lines


# Add here suffixes that are used to identify models, separated by |
_SCREAMING_SNAKE_CASE					:				Tuple					    =						'''Model|Encoder|Decoder|ForConditionalGeneration'''
# Regexes that match TF/Flax/PT model names.
_SCREAMING_SNAKE_CASE					:				Tuple					    =						re.compile(R'''TF(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)''')
_SCREAMING_SNAKE_CASE					:				int					    =						re.compile(R'''Flax(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)''')
# Will match any TF or Flax model too so need to be in an else branch afterthe two previous regexes.
_SCREAMING_SNAKE_CASE					:				Union[str, Any]					    =						re.compile(R'''(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)''')


# This is to make sure the transformers module imported is the one in the repo.
_SCREAMING_SNAKE_CASE					:				Optional[int]					    =						direct_transformers_import(TRANSFORMERS_PATH)







def      lowerCamelCase__       (							_lowerCamelCase       :		Any )      ->				List[str]:
  lowerCamelCase_         =						re.finditer('.+?(?:(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|$)'		,				_lowerCamelCase )
  return [m.group(0 ) for m in matches]







def      lowerCamelCase__       (							_lowerCamelCase       :		Dict		,				_lowerCamelCase       :		Any )      ->				Optional[Any]:
  lowerCamelCase_         =						2 if text == '✅' or text == '❌' else len(_lowerCamelCase )
  lowerCamelCase_         =						(width - text_length) // 2
  lowerCamelCase_         =						width - text_length - left_indent
  return " " * left_indent + text + " " * right_indent







def      lowerCamelCase__       (							)      ->				Optional[Any]:
  lowerCamelCase_         =						transformers_module.models.auto.configuration_auto.CONFIG_MAPPING_NAMES
  lowerCamelCase_         =						{
      name: config_maping_names[code]
      for code, name in transformers_module.MODEL_NAMES_MAPPING.items()
      if code in config_maping_names
  }
  lowerCamelCase_         =						{name: config.replace('Config'		,				'' ) for name, config in model_name_to_config.items()}

  # Dictionaries flagging if each model prefix has a slow/fast tokenizer, backend in PT/TF/Flax.
  lowerCamelCase_         =						collections.defaultdict(_lowerCamelCase )
  lowerCamelCase_         =						collections.defaultdict(_lowerCamelCase )
  lowerCamelCase_         =						collections.defaultdict(_lowerCamelCase )
  lowerCamelCase_         =						collections.defaultdict(_lowerCamelCase )
  lowerCamelCase_         =						collections.defaultdict(_lowerCamelCase )

  # Let's lookup through all transformers object (once).
  for attr_name in dir(_lowerCamelCase ):
    lowerCamelCase_         =						None
    if attr_name.endswith('Tokenizer' ):
      lowerCamelCase_         =						slow_tokenizers
      lowerCamelCase_         =						attr_name[:-9]
    elif attr_name.endswith('TokenizerFast' ):
      lowerCamelCase_         =						fast_tokenizers
      lowerCamelCase_         =						attr_name[:-13]
    elif _re_tf_models.match(_lowerCamelCase ) is not None:
      lowerCamelCase_         =						tf_models
      lowerCamelCase_         =						_re_tf_models.match(_lowerCamelCase ).groups()[0]
    elif _re_flax_models.match(_lowerCamelCase ) is not None:
      lowerCamelCase_         =						flax_models
      lowerCamelCase_         =						_re_flax_models.match(_lowerCamelCase ).groups()[0]
    elif _re_pt_models.match(_lowerCamelCase ) is not None:
      lowerCamelCase_         =						pt_models
      lowerCamelCase_         =						_re_pt_models.match(_lowerCamelCase ).groups()[0]

    if lookup_dict is not None:
      while len(_lowerCamelCase ) > 0:
        if attr_name in model_name_to_prefix.values():
          lowerCamelCase_         =						True
          break
        # Try again after removing the last word in the name
        lowerCamelCase_         =						''.join(camel_case_split(_lowerCamelCase )[:-1] )

    # Let's build that table!
  lowerCamelCase_         =						list(model_name_to_config.keys() )
  model_names.sort(key=str.lower )
  lowerCamelCase_         =						['Model', 'Tokenizer slow', 'Tokenizer fast', 'PyTorch support', 'TensorFlow support', 'Flax Support']
  # We'll need widths to properly display everything in the center (+2 is to leave one extra space on each side).
  lowerCamelCase_         =						[len(_lowerCamelCase ) + 2 for c in columns]
  lowerCamelCase_         =						max([len(_lowerCamelCase ) for name in model_names] ) + 2

  # Build the table per se
  lowerCamelCase_         =						'|' + '|'.join([_center_text(_lowerCamelCase		,				_lowerCamelCase ) for c, w in zip(_lowerCamelCase		,				_lowerCamelCase )] ) + '|\n'
  # Use ":-----:" format to center-aligned table cell texts
  table += "|" + "|".join([':' + '-' * (w - 2) + ':' for w in widths] ) + "|\n"

  lowerCamelCase_         =						{True: '✅', False: '❌'}
  for name in model_names:
    lowerCamelCase_         =						model_name_to_prefix[name]
    lowerCamelCase_         =						[
        name,
        check[slow_tokenizers[prefix]],
        check[fast_tokenizers[prefix]],
        check[pt_models[prefix]],
        check[tf_models[prefix]],
        check[flax_models[prefix]],
    ]
    table += "|" + "|".join([_center_text(_lowerCamelCase		,				_lowerCamelCase ) for l, w in zip(_lowerCamelCase		,				_lowerCamelCase )] ) + "|\n"
  return table







def      lowerCamelCase__       (							_lowerCamelCase       :		Optional[int]=False )      ->				Any:
  lowerCamelCase_     ,	lowerCamelCase_     ,	lowerCamelCase_     ,	lowerCamelCase_         =						_find_text_in_file(
      filename=os.path.join(_lowerCamelCase		,				'index.md' )		,				start_prompt='<!--This table is updated automatically from the auto modules'		,				end_prompt='<!-- End table-->'		,				)
  lowerCamelCase_         =						get_model_table_from_auto_modules()

  if current_table != new_table:
    if overwrite:
      with open(os.path.join(_lowerCamelCase		,				'index.md' )		,				'w'		,				encoding='utf-8'		,				newline='\n' ) as f:
        f.writelines(lines[:start_index] + [new_table] + lines[end_index:] )
    else:
      raise ValueError(
          'The model table in the `index.md` has not been updated. Run `make fix-copies` to fix this.' )


if __name__ == "__main__":
  _SCREAMING_SNAKE_CASE					:				Union[str, Any]					    =						argparse.ArgumentParser()
  parser.add_argument('''--fix_and_overwrite''', action='''store_true''', help='''Whether to fix inconsistencies.''')
  _SCREAMING_SNAKE_CASE					:				Any					    =						parser.parse_args()

  check_model_table(args.fix_and_overwrite)












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


_SCREAMING_SNAKE_CASE					:				Union[str, Any]					    =						['''\nclass''', '''\ndef''', '''\n#''', '''\n@''', '''\nprint''', '''\nif''']






class      a						(					__snake_case		):






    def __init__(       self     :							Dict       ,						__SCREAMING_SNAKE_CASE     :							Any       ,						__SCREAMING_SNAKE_CASE     :							Optional[Any]       ,						__SCREAMING_SNAKE_CASE     :							Any=None       ,						__SCREAMING_SNAKE_CASE     :							Tuple=1					)     ->		str:
      lowerCamelCase_         =						tokenizer
      lowerCamelCase_         =						dataset
      lowerCamelCase_         =						len(__SCREAMING_SNAKE_CASE					) if n_tasks is None else n_tasks
      lowerCamelCase_         =						n_copies






    def __iter__(       self     :							Dict					)     ->		Any:
      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(__SCREAMING_SNAKE_CASE       ,						padding=__SCREAMING_SNAKE_CASE       ,						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						(					__snake_case		):






    def __init__(       self     :							Union[str, Any]       ,						__SCREAMING_SNAKE_CASE     :							str       ,						__SCREAMING_SNAKE_CASE     :							int       ,						__SCREAMING_SNAKE_CASE     :							Any					)     ->		List[Any]:
      lowerCamelCase_         =						start_length
      lowerCamelCase_         =						eof_strings
      lowerCamelCase_         =						tokenizer






    def __call__(       self     :							Union[str, Any]       ,						__SCREAMING_SNAKE_CASE     :							Optional[int]       ,						__SCREAMING_SNAKE_CASE     :							Dict       ,						**__SCREAMING_SNAKE_CASE     :							Union[str, Any]					)     ->		Optional[Any]:
      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(__SCREAMING_SNAKE_CASE					)







def      lowerCamelCase__       (							_lowerCamelCase       :		List[Any] )      ->				Tuple:
  lowerCamelCase_         =						re.split('(%s)' % '|'.join(_lowerCamelCase )		,				_lowerCamelCase )
  # last string should be ""
  return "".join(string_list[:-2] )







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

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







def      lowerCamelCase__       (							)      ->				Tuple:
  # Setup configuration
  lowerCamelCase_         =						HfArgumentParser(_lowerCamelCase )
  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=_lowerCamelCase )

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

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

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

  if accelerator.is_main_process:
    lowerCamelCase_         =						[]

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




