


"""simple docstring"""


from __future__ import annotations

import unittest

import numpy as np

from transformers import LayoutLMConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
  import tensorflow as tf

  from transformers.models.layoutlm.modeling_tf_layoutlm import (
      TF_LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFLayoutLMForMaskedLM,
      TFLayoutLMForQuestionAnswering,
      TFLayoutLMForSequenceClassification,
      TFLayoutLMForTokenClassification,
      TFLayoutLMModel,
  )



class   _a      :
    def __init__( self							:					Optional[int]							,					SCREAMING_SNAKE_CASE__							:					Dict							,					SCREAMING_SNAKE_CASE__							:					str=13							,					SCREAMING_SNAKE_CASE__							:					List[str]=7							,					SCREAMING_SNAKE_CASE__							:					List[Any]=True							,					SCREAMING_SNAKE_CASE__							:					Any=True							,					SCREAMING_SNAKE_CASE__							:					List[str]=True							,					SCREAMING_SNAKE_CASE__							:					Dict=True							,					SCREAMING_SNAKE_CASE__							:					int=99							,					SCREAMING_SNAKE_CASE__							:					Any=32							,					SCREAMING_SNAKE_CASE__							:					int=2							,					SCREAMING_SNAKE_CASE__							:					str=4							,					SCREAMING_SNAKE_CASE__							:					Optional[int]=37							,					SCREAMING_SNAKE_CASE__							:					Any="gelu"							,					SCREAMING_SNAKE_CASE__							:					Tuple=0.1							,					SCREAMING_SNAKE_CASE__							:					Union[str, Any]=0.1							,					SCREAMING_SNAKE_CASE__							:					Optional[Any]=5_12							,					SCREAMING_SNAKE_CASE__							:					Any=16							,					SCREAMING_SNAKE_CASE__							:					int=2							,					SCREAMING_SNAKE_CASE__							:					Any=0.02							,					SCREAMING_SNAKE_CASE__							:					List[str]=3							,					SCREAMING_SNAKE_CASE__							:					int=4							,					SCREAMING_SNAKE_CASE__							:					Tuple=None							,					SCREAMING_SNAKE_CASE__							:					List[Any]=10_00							,					):
          lowerCamelCase__           =		parent
          lowerCamelCase__           =		batch_size
          lowerCamelCase__           =		seq_length
          lowerCamelCase__           =		is_training
          lowerCamelCase__           =		use_input_mask
          lowerCamelCase__           =		use_token_type_ids
          lowerCamelCase__           =		use_labels
          lowerCamelCase__           =		vocab_size
          lowerCamelCase__           =		hidden_size
          lowerCamelCase__           =		num_hidden_layers
          lowerCamelCase__           =		num_attention_heads
          lowerCamelCase__           =		intermediate_size
          lowerCamelCase__           =		hidden_act
          lowerCamelCase__           =		hidden_dropout_prob
          lowerCamelCase__           =		attention_probs_dropout_prob
          lowerCamelCase__           =		max_position_embeddings
          lowerCamelCase__           =		type_vocab_size
          lowerCamelCase__           =		type_sequence_label_size
          lowerCamelCase__           =		initializer_range
          lowerCamelCase__           =		num_labels
          lowerCamelCase__           =		num_choices
          lowerCamelCase__           =		scope
          lowerCamelCase__           =		range_bbox
    def      _UpperCamelCase       ( self							:					Any				):
          lowerCamelCase__           =		ids_tensor([self.batch_size, self.seq_length]							,					self.vocab_size				)

          # convert bbox to numpy since TF does not support item assignment
          lowerCamelCase__           =		ids_tensor([self.batch_size, self.seq_length, 4]							,					self.range_bbox				).numpy()
          # Ensure that bbox is legal
          for i in range(bbox.shape[0]				):
                for j in range(bbox.shape[1]				):
                      if bbox[i, j, 3] < bbox[i, j, 1]:
                            lowerCamelCase__           =		bbox[i, j, 3]
                            lowerCamelCase__           =		bbox[i, j, 1]
                            lowerCamelCase__           =		t
                      if bbox[i, j, 2] < bbox[i, j, 0]:
                            lowerCamelCase__           =		bbox[i, j, 2]
                            lowerCamelCase__           =		bbox[i, j, 0]
                            lowerCamelCase__           =		t
          lowerCamelCase__           =		tf.convert_to_tensor(SCREAMING_SNAKE_CASE__				)

          lowerCamelCase__           =		None
          if self.use_input_mask:
                lowerCamelCase__           =		random_attention_mask([self.batch_size, self.seq_length]				)

          lowerCamelCase__           =		None
          if self.use_token_type_ids:
                lowerCamelCase__           =		ids_tensor([self.batch_size, self.seq_length]							,					self.type_vocab_size				)

          lowerCamelCase__           =		None
          lowerCamelCase__           =		None
          lowerCamelCase__           =		None
          if self.use_labels:
                lowerCamelCase__           =		ids_tensor([self.batch_size]							,					self.type_sequence_label_size				)
                lowerCamelCase__           =		ids_tensor([self.batch_size, self.seq_length]							,					self.num_labels				)
                lowerCamelCase__           =		ids_tensor([self.batch_size]							,					self.num_choices				)

          lowerCamelCase__           =		LayoutLMConfig(
              vocab_size=self.vocab_size							,					hidden_size=self.hidden_size							,					num_hidden_layers=self.num_hidden_layers							,					num_attention_heads=self.num_attention_heads							,					intermediate_size=self.intermediate_size							,					hidden_act=self.hidden_act							,					hidden_dropout_prob=self.hidden_dropout_prob							,					attention_probs_dropout_prob=self.attention_probs_dropout_prob							,					max_position_embeddings=self.max_position_embeddings							,					type_vocab_size=self.type_vocab_size							,					initializer_range=self.initializer_range							,					)

          return config, input_ids, bbox, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
    def      _UpperCamelCase       ( self							:					Union[str, Any]							,					SCREAMING_SNAKE_CASE__							:					Dict							,					SCREAMING_SNAKE_CASE__							:					Any							,					SCREAMING_SNAKE_CASE__							:					Dict							,					SCREAMING_SNAKE_CASE__							:					Optional[Any]							,					SCREAMING_SNAKE_CASE__							:					List[str]							,					SCREAMING_SNAKE_CASE__							:					Optional[Any]							,					SCREAMING_SNAKE_CASE__							:					Tuple							,					SCREAMING_SNAKE_CASE__							:					Optional[Any]				):
          lowerCamelCase__           =		TFLayoutLMModel(config=SCREAMING_SNAKE_CASE__				)

          lowerCamelCase__           =		model(SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__							,					attention_mask=SCREAMING_SNAKE_CASE__							,					token_type_ids=SCREAMING_SNAKE_CASE__				)
          lowerCamelCase__           =		model(SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__							,					token_type_ids=SCREAMING_SNAKE_CASE__				)
          lowerCamelCase__           =		model(SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__				)
          self.parent.assertEqual(result.last_hidden_state.shape							,					(self.batch_size, self.seq_length, self.hidden_size)				)
          self.parent.assertEqual(result.pooler_output.shape							,					(self.batch_size, self.hidden_size)				)
    def      _UpperCamelCase       ( self							:					List[str]							,					SCREAMING_SNAKE_CASE__							:					Dict							,					SCREAMING_SNAKE_CASE__							:					int							,					SCREAMING_SNAKE_CASE__							:					Any							,					SCREAMING_SNAKE_CASE__							:					Any							,					SCREAMING_SNAKE_CASE__							:					Dict							,					SCREAMING_SNAKE_CASE__							:					Optional[Any]							,					SCREAMING_SNAKE_CASE__							:					int							,					SCREAMING_SNAKE_CASE__							:					List[str]				):
          lowerCamelCase__           =		TFLayoutLMForMaskedLM(config=SCREAMING_SNAKE_CASE__				)

          lowerCamelCase__           =		model(SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__							,					attention_mask=SCREAMING_SNAKE_CASE__							,					token_type_ids=SCREAMING_SNAKE_CASE__							,					labels=SCREAMING_SNAKE_CASE__				)
          self.parent.assertEqual(result.logits.shape							,					(self.batch_size, self.seq_length, self.vocab_size)				)
    def      _UpperCamelCase       ( self							:					Union[str, Any]							,					SCREAMING_SNAKE_CASE__							:					Union[str, Any]							,					SCREAMING_SNAKE_CASE__							:					Any							,					SCREAMING_SNAKE_CASE__							:					int							,					SCREAMING_SNAKE_CASE__							:					List[str]							,					SCREAMING_SNAKE_CASE__							:					int							,					SCREAMING_SNAKE_CASE__							:					List[str]							,					SCREAMING_SNAKE_CASE__							:					Optional[int]							,					SCREAMING_SNAKE_CASE__							:					Union[str, Any]				):
          lowerCamelCase__           =		self.num_labels
          lowerCamelCase__           =		TFLayoutLMForSequenceClassification(config=SCREAMING_SNAKE_CASE__				)

          lowerCamelCase__           =		model(SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__							,					attention_mask=SCREAMING_SNAKE_CASE__							,					token_type_ids=SCREAMING_SNAKE_CASE__				)
          self.parent.assertEqual(result.logits.shape							,					(self.batch_size, self.num_labels)				)
    def      _UpperCamelCase       ( self							:					Dict							,					SCREAMING_SNAKE_CASE__							:					str							,					SCREAMING_SNAKE_CASE__							:					Dict							,					SCREAMING_SNAKE_CASE__							:					List[Any]							,					SCREAMING_SNAKE_CASE__							:					List[Any]							,					SCREAMING_SNAKE_CASE__							:					int							,					SCREAMING_SNAKE_CASE__							:					int							,					SCREAMING_SNAKE_CASE__							:					Union[str, Any]							,					SCREAMING_SNAKE_CASE__							:					Dict				):
          lowerCamelCase__           =		self.num_labels
          lowerCamelCase__           =		TFLayoutLMForTokenClassification(config=SCREAMING_SNAKE_CASE__				)

          lowerCamelCase__           =		model(SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__							,					attention_mask=SCREAMING_SNAKE_CASE__							,					token_type_ids=SCREAMING_SNAKE_CASE__							,					labels=SCREAMING_SNAKE_CASE__				)
          self.parent.assertEqual(result.logits.shape							,					(self.batch_size, self.seq_length, self.num_labels)				)
    def      _UpperCamelCase       ( self							:					str							,					SCREAMING_SNAKE_CASE__							:					List[Any]							,					SCREAMING_SNAKE_CASE__							:					Tuple							,					SCREAMING_SNAKE_CASE__							:					Tuple							,					SCREAMING_SNAKE_CASE__							:					Any							,					SCREAMING_SNAKE_CASE__							:					Union[str, Any]							,					SCREAMING_SNAKE_CASE__							:					Dict							,					SCREAMING_SNAKE_CASE__							:					Union[str, Any]							,					SCREAMING_SNAKE_CASE__							:					List[str]				):
          lowerCamelCase__           =		TFLayoutLMForQuestionAnswering(config=SCREAMING_SNAKE_CASE__				)

          lowerCamelCase__           =		model(SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__							,					attention_mask=SCREAMING_SNAKE_CASE__							,					token_type_ids=SCREAMING_SNAKE_CASE__				)
          self.parent.assertEqual(result.start_logits.shape							,					(self.batch_size, self.seq_length)				)
          self.parent.assertEqual(result.end_logits.shape							,					(self.batch_size, self.seq_length)				)


    def      _UpperCamelCase       ( self							:					str				):
          lowerCamelCase__           =		self.prepare_config_and_inputs()
          (
              (
              lowerCamelCase__  
          )    ,  (
              lowerCamelCase__  
          )    ,  (
              lowerCamelCase__  
          )    ,  (
              lowerCamelCase__  
          )    ,  (
              lowerCamelCase__  
          )    ,  (
              lowerCamelCase__  
          )    ,  (
              lowerCamelCase__  
          )    ,  (
              lowerCamelCase__  
          )    ,    
          )           =		config_and_inputs
          lowerCamelCase__           =		{
              'input_ids': input_ids,
              'bbox': bbox,
              'token_type_ids': token_type_ids,
              'attention_mask': input_mask,
          }
          return config, inputs_dict





@require_tf
class   _a      (      SCREAMING_SNAKE_CASE_		,		SCREAMING_SNAKE_CASE_		,		unittest.TestCase ):
    a_    :				Optional[Any]             =    (
        (
            TFLayoutLMModel,
            TFLayoutLMForMaskedLM,
            TFLayoutLMForTokenClassification,
            TFLayoutLMForSequenceClassification,
            TFLayoutLMForQuestionAnswering,
        )
        if is_tf_available()
        else ()
    )
    a_    :				Tuple             =    (
        {
            'feature-extraction': TFLayoutLMModel,
            'fill-mask': TFLayoutLMForMaskedLM,
            'text-classification': TFLayoutLMForSequenceClassification,
            'token-classification': TFLayoutLMForTokenClassification,
            'zero-shot': TFLayoutLMForSequenceClassification,
        }
        if is_tf_available()
        else {}
    )
    a_    :				Optional[Any]             =    False
    a_    :				int             =    True
    a_    :				Optional[Any]             =    10
    def      _UpperCamelCase       ( self							:					Tuple				):
          lowerCamelCase__           =		TFLayoutLMModelTester(self				)
          lowerCamelCase__           =		ConfigTester(self							,					config_class=SCREAMING_SNAKE_CASE__							,					hidden_size=37				)
    def      _UpperCamelCase       ( self							:					Optional[int]				):
          self.config_tester.run_common_tests()
    def      _UpperCamelCase       ( self							:					Any				):
          lowerCamelCase__           =		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE__				)
    def      _UpperCamelCase       ( self							:					List[Any]				):
          lowerCamelCase__           =		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE__				)
    def      _UpperCamelCase       ( self							:					Optional[int]				):
          lowerCamelCase__           =		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_sequence_classification(*SCREAMING_SNAKE_CASE__				)
    def      _UpperCamelCase       ( self							:					Optional[Any]				):
          lowerCamelCase__           =		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE__				)
    def      _UpperCamelCase       ( self							:					Any				):
          lowerCamelCase__           =		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_question_answering(*SCREAMING_SNAKE_CASE__				)
    @slow
    def      _UpperCamelCase       ( self							:					List[Any]				):
          for model_name in TF_LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                lowerCamelCase__           =		TFLayoutLMModel.from_pretrained(SCREAMING_SNAKE_CASE__				)
                self.assertIsNotNone(SCREAMING_SNAKE_CASE__				)


    @unittest.skip('Onnx compliancy broke with TF 2.10'				)
    def      _UpperCamelCase       ( self							:					str				):
          pass





def   snake_case						(				)->		Union[str, Any]:
      '''simple docstring'''


      lowerCamelCase__           =		tf.convert_to_tensor([[101,1019,1014,1016,1037,12849,4747,1004,14246,2278,5439,4524,5002,2930,2193,2930,4341,3208,1005,1055,2171,2848,11300,3531,102],[101,4070,4034,7020,1024,3058,1015,1013,2861,1013,6070,19274,2772,6205,27814,16147,16147,4343,2047,10283,10969,14389,1012,2338,102]]							)  # noqa: E231
      lowerCamelCase__           =		tf.convert_to_tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],]							)  # noqa: E231
      lowerCamelCase__           =		tf.convert_to_tensor([[[0,0,0,0],[423,237,440,251],[427,272,441,287],[419,115,437,129],[961,885,992,912],[256,38,330,58],[256,38,330,58],[336,42,353,57],[360,39,401,56],[360,39,401,56],[411,39,471,59],[479,41,528,59],[533,39,630,60],[67,113,134,131],[141,115,209,132],[68,149,133,166],[141,149,187,164],[195,148,287,165],[195,148,287,165],[195,148,287,165],[295,148,349,165],[441,149,492,166],[497,149,546,164],[64,201,125,218],[1000,1000,1000,1000]],[[0,0,0,0],[662,150,754,166],[665,199,742,211],[519,213,554,228],[519,213,554,228],[134,433,187,454],[130,467,204,480],[130,467,204,480],[130,467,204,480],[130,467,204,480],[130,467,204,480],[314,469,376,482],[504,684,582,706],[941,825,973,900],[941,825,973,900],[941,825,973,900],[941,825,973,900],[610,749,652,765],[130,659,168,672],[176,657,237,672],[238,657,312,672],[443,653,628,672],[443,653,628,672],[716,301,825,317],[1000,1000,1000,1000]]]							)  # noqa: E231
      lowerCamelCase__           =		tf.convert_to_tensor([[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]]							)  # noqa: E231
      # these are sequence labels (i.e. at the token level)
      lowerCamelCase__           =		tf.convert_to_tensor([[-100,10,10,10,9,1,-100,7,7,-100,7,7,4,2,5,2,8,8,-100,-100,5,0,3,2,-100],[-100,12,12,12,-100,12,10,-100,-100,-100,-100,10,12,9,-100,-100,-100,10,10,10,9,12,-100,10,-100]]							)  # noqa: E231
      # fmt: on

      return input_ids, attention_mask, bbox, token_type_ids, labels






@require_tf
class   _a      (      unittest.TestCase ):
    @slow
    def      _UpperCamelCase       ( self							:					Union[str, Any]				):
          lowerCamelCase__           =		TFLayoutLMModel.from_pretrained('microsoft/layoutlm-base-uncased'				)

          lowerCamelCase__    ,  lowerCamelCase__    ,  lowerCamelCase__    ,  lowerCamelCase__    ,  lowerCamelCase__           =		prepare_layoutlm_batch_inputs()

          # forward pass
          lowerCamelCase__           =		model(input_ids=SCREAMING_SNAKE_CASE__							,					bbox=SCREAMING_SNAKE_CASE__							,					attention_mask=SCREAMING_SNAKE_CASE__							,					token_type_ids=SCREAMING_SNAKE_CASE__				)

          # test the sequence output on [0, :3, :3]
          lowerCamelCase__           =		tf.convert_to_tensor(
              [[0.17_85, -0.19_47, -0.04_25], [-0.32_54, -0.28_07, 0.25_53], [-0.53_91, -0.33_22, 0.33_64]]							,					)

          self.assertTrue(np.allclose(outputs.last_hidden_state[0, :3, :3]							,					SCREAMING_SNAKE_CASE__							,					atol=1e-3				)				)

          # test the pooled output on [1, :3]
          lowerCamelCase__           =		tf.convert_to_tensor([-0.65_80, -0.02_14, 0.85_52]				)

          self.assertTrue(np.allclose(outputs.pooler_output[1, :3]							,					SCREAMING_SNAKE_CASE__							,					atol=1e-3				)				)
    @slow
    def      _UpperCamelCase       ( self							:					str				):
          # initialize model with randomly initialized sequence classification head
          lowerCamelCase__           =		TFLayoutLMForSequenceClassification.from_pretrained('microsoft/layoutlm-base-uncased'							,					num_labels=2				)

          lowerCamelCase__    ,  lowerCamelCase__    ,  lowerCamelCase__    ,  lowerCamelCase__    ,  lowerCamelCase__           =		prepare_layoutlm_batch_inputs()

          # forward pass
          lowerCamelCase__           =		model(
              input_ids=SCREAMING_SNAKE_CASE__							,					bbox=SCREAMING_SNAKE_CASE__							,					attention_mask=SCREAMING_SNAKE_CASE__							,					token_type_ids=SCREAMING_SNAKE_CASE__							,					labels=tf.convert_to_tensor([1, 1]				)							,					)

          # test whether we get a loss as a scalar
          lowerCamelCase__           =		outputs.loss
          lowerCamelCase__           =		(2,)
          self.assertEqual(loss.shape							,					SCREAMING_SNAKE_CASE__				)

          # test the shape of the logits
          lowerCamelCase__           =		outputs.logits
          lowerCamelCase__           =		(2, 2)
          self.assertEqual(logits.shape							,					SCREAMING_SNAKE_CASE__				)
    @slow
    def      _UpperCamelCase       ( self							:					Union[str, Any]				):
          # initialize model with randomly initialized token classification head
          lowerCamelCase__           =		TFLayoutLMForTokenClassification.from_pretrained('microsoft/layoutlm-base-uncased'							,					num_labels=13				)

          lowerCamelCase__    ,  lowerCamelCase__    ,  lowerCamelCase__    ,  lowerCamelCase__    ,  lowerCamelCase__           =		prepare_layoutlm_batch_inputs()

          # forward pass
          lowerCamelCase__           =		model(
              input_ids=SCREAMING_SNAKE_CASE__							,					bbox=SCREAMING_SNAKE_CASE__							,					attention_mask=SCREAMING_SNAKE_CASE__							,					token_type_ids=SCREAMING_SNAKE_CASE__							,					labels=SCREAMING_SNAKE_CASE__				)

          # test the shape of the logits
          lowerCamelCase__           =		outputs.logits
          lowerCamelCase__           =		tf.convert_to_tensor((2, 25, 13)				)
          self.assertEqual(logits.shape							,					SCREAMING_SNAKE_CASE__				)


    @slow
    def      _UpperCamelCase       ( self							:					Optional[Any]				):
          # initialize model with randomly initialized token classification head
          lowerCamelCase__           =		TFLayoutLMForQuestionAnswering.from_pretrained('microsoft/layoutlm-base-uncased'				)

          lowerCamelCase__    ,  lowerCamelCase__    ,  lowerCamelCase__    ,  lowerCamelCase__    ,  lowerCamelCase__           =		prepare_layoutlm_batch_inputs()

          # forward pass
          lowerCamelCase__           =		model(input_ids=SCREAMING_SNAKE_CASE__							,					bbox=SCREAMING_SNAKE_CASE__							,					attention_mask=SCREAMING_SNAKE_CASE__							,					token_type_ids=SCREAMING_SNAKE_CASE__				)

          # test the shape of the logits
          lowerCamelCase__           =		tf.convert_to_tensor((2, 25)				)
          self.assertEqual(outputs.start_logits.shape							,					SCREAMING_SNAKE_CASE__				)
          self.assertEqual(outputs.end_logits.shape							,					SCREAMING_SNAKE_CASE__				)







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


_snake_case          =							["\nclass", "\ndef", "\n#", "\n@", "\nprint", "\nif"]



class   _a      (      SCREAMING_SNAKE_CASE_ ):
    def __init__( self							:					List[str]							,					SCREAMING_SNAKE_CASE__							:					Any							,					SCREAMING_SNAKE_CASE__							:					List[Any]							,					SCREAMING_SNAKE_CASE__							:					int=None							,					SCREAMING_SNAKE_CASE__							:					Optional[int]=1				):
          lowerCamelCase__           =		tokenizer
          lowerCamelCase__           =		dataset
          lowerCamelCase__           =		len(SCREAMING_SNAKE_CASE__				) if n_tasks is None else n_tasks
          lowerCamelCase__           =		n_copies


    def __iter__( self							:					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      (      SCREAMING_SNAKE_CASE_ ):
    def __init__( self							:					Dict							,					SCREAMING_SNAKE_CASE__							:					List[str]							,					SCREAMING_SNAKE_CASE__							:					List[Any]							,					SCREAMING_SNAKE_CASE__							:					Dict				):
          lowerCamelCase__           =		start_length
          lowerCamelCase__           =		eof_strings
          lowerCamelCase__           =		tokenizer


    def __call__( self							:					Any							,					SCREAMING_SNAKE_CASE__							:					Any							,					SCREAMING_SNAKE_CASE__							:					List[str]							,					**SCREAMING_SNAKE_CASE__							:					Dict				):
          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   snake_case						(				_a:     List[Any]							)->		Dict:
      '''simple docstring'''


      lowerCamelCase__           =		re.split('(%s)' % '|'.join(_a							)		,   _a							)
      # last string should be ""
      return "".join(string_list[:-2]							)



def   snake_case						(				_a:     List[Any]		,   _a:     Optional[int]		,   _a:     str		,   _a:     Union[str, Any]		,   _a:     Dict		,   _a:     Optional[int]=20		,   **_a:     Optional[int]							)->		List[str]:
      '''simple docstring'''


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

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



def   snake_case						(				)->		Union[str, Any]:
      '''simple docstring'''


      lowerCamelCase__           =		HfArgumentParser(_a							)
      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=_a							)

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

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

      lowerCamelCase__           =		complete_code(
          _a		,   _a		,   _a		,   _a		,   n_tasks=_a		,   batch_size=args.batch_size		,   **_a		,   )

      if accelerator.is_main_process:
            lowerCamelCase__           =		[]

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


# 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()



