



import logging

import torch
from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss

from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.bert.modeling_bert import (
    BERT_INPUTS_DOCSTRING,
    BERT_START_DOCSTRING,
    BertEncoder,
    BertModel,
    BertPreTrainedModel,
)


lowerCamelCase_  :	int 				=					logging.getLogger(__name__)

class     _UpperCamelCase			(   __UpperCamelCase					):






			'''simple docstring'''

			def  lowerCAmelCase__		(    self	:    Union[str, Any]  ,		snake_case_	:    Optional[int]  ,		snake_case_	:    Optional[Any]  ,		snake_case_	:    Tuple=None  ,		snake_case_	:    Optional[int]=None ):
							UpperCamelCase_:							List[str]		=  self.layer[current_layer](snake_case_  ,		snake_case_  ,		head_mask[current_layer] )

							UpperCamelCase_:							Optional[int]		=  layer_outputs[0]

							return hidden_states




@add_start_docstrings(
    """The bare Bert Model transformer with PABEE outputting raw hidden-states without any specific head on top."""       ,    __UpperCamelCase       ,    )
class     _UpperCamelCase			(   __UpperCamelCase					):






			'''simple docstring'''

			def __init__(    self	:    Dict  ,		snake_case_	:    Union[str, Any] ):
							super().__init__(snake_case_ )

							UpperCamelCase_:							int		=  BertEncoderWithPabee(snake_case_ )

							self.init_weights()
							UpperCamelCase_:							Union[str, Any]		=  0
							UpperCamelCase_:							int		=  0
							UpperCamelCase_:							str		=  0

							UpperCamelCase_:							List[str]		=  0

			def  lowerCAmelCase__		(    self	:    List[Any]  ,		snake_case_	:    Tuple ):
							UpperCamelCase_:							List[Any]		=  threshold

			def  lowerCAmelCase__		(    self	:    Any  ,		snake_case_	:    str ):
							UpperCamelCase_:							Any		=  patience

			def  lowerCAmelCase__		(    self	:    List[Any] ):
							UpperCamelCase_:							Dict		=  0
							UpperCamelCase_:							List[Any]		=  0

			def  lowerCAmelCase__		(    self	:    int ):
							UpperCamelCase_:							Optional[int]		=  self.inference_layers_num / self.inference_instances_num
							UpperCamelCase_:							Optional[Any]		=  (
							    f'''*** Patience = {self.patience} Avg. Inference Layers = {avg_inf_layers:.2f} Speed Up ='''
							    f''' {1 - avg_inf_layers / self.config.num_hidden_layers:.2f} ***'''
							)
							print(snake_case_ )




			@add_start_docstrings_to_model_forward(snake_case_ )
			def  lowerCAmelCase__		(    self	:    Any  ,		snake_case_	:    str=None  ,		snake_case_	:    Tuple=None  ,		snake_case_	:    str=None  ,		snake_case_	:    int=None  ,		snake_case_	:    List[Any]=None  ,		snake_case_	:    Optional[Any]=None  ,		snake_case_	:    List[str]=None  ,		snake_case_	:    Optional[int]=None  ,		snake_case_	:    Tuple=None  ,		snake_case_	:    str=None  ,		snake_case_	:    Optional[int]=False  ,		):

							if input_ids is not None and inputs_embeds is not None:
											raise ValueError("""You cannot specify both input_ids and inputs_embeds at the same time""" )
							elif input_ids is not None:
											UpperCamelCase_:							List[Any]		=  input_ids.size()
							elif inputs_embeds is not None:
											UpperCamelCase_:							List[Any]		=  inputs_embeds.size()[:-1]
							else:
											raise ValueError("""You have to specify either input_ids or inputs_embeds""" )

							UpperCamelCase_:							Union[str, Any]		=  input_ids.device if input_ids is not None else inputs_embeds.device

							if attention_mask is None:
											UpperCamelCase_:							str		=  torch.ones(snake_case_  ,		device=snake_case_ )
							if token_type_ids is None:
											UpperCamelCase_:							int		=  torch.zeros(snake_case_  ,		dtype=torch.long  ,		device=snake_case_ )

							# We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
							# ourselves in which case we just need to make it broadcastable to all heads.
							UpperCamelCase_:							torch.Tensor		=  self.get_extended_attention_mask(snake_case_  ,		snake_case_  ,		snake_case_ )

							# If a 2D ou 3D attention mask is provided for the cross-attention
							# we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
							if self.config.is_decoder and encoder_hidden_states is not None:
											UpperCamelCase_:							Union[str, Any]		=  encoder_hidden_states.size()
											UpperCamelCase_:							Any		=  (encoder_batch_size, encoder_sequence_length)
											if encoder_attention_mask is None:
															UpperCamelCase_:							Optional[int]		=  torch.ones(snake_case_  ,		device=snake_case_ )
											UpperCamelCase_:							str		=  self.invert_attention_mask(snake_case_ )
							else:
											UpperCamelCase_:							int		=  None

							# Prepare head mask if needed
							# 1.0 in head_mask indicate we keep the head
							# attention_probs has shape bsz x n_heads x N x N
							# input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
							# and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
							UpperCamelCase_:							Optional[Any]		=  self.get_head_mask(snake_case_  ,		self.config.num_hidden_layers )

							UpperCamelCase_:							Optional[Any]		=  self.embeddings(
							    input_ids=snake_case_  ,		position_ids=snake_case_  ,		token_type_ids=snake_case_  ,		inputs_embeds=snake_case_ )
							UpperCamelCase_:							Tuple		=  embedding_output

							if self.training:
											UpperCamelCase_:							Optional[Any]		=  []
											for i in range(self.config.num_hidden_layers ):
															UpperCamelCase_:							Dict		=  self.encoder.adaptive_forward(
															    snake_case_  ,		current_layer=snake_case_  ,		attention_mask=snake_case_  ,		head_mask=snake_case_ )

															UpperCamelCase_:							Optional[Any]		=  self.pooler(snake_case_ )
															UpperCamelCase_:							Any		=  output_layers[i](output_dropout(snake_case_ ) )
															res.append(snake_case_ )
							elif self.patience == 0:  # Use all layers for inference
											UpperCamelCase_:							List[Any]		=  self.encoder(
											    snake_case_  ,		attention_mask=snake_case_  ,		head_mask=snake_case_  ,		encoder_hidden_states=snake_case_  ,		encoder_attention_mask=snake_case_  ,		)
											UpperCamelCase_:							List[Any]		=  self.pooler(encoder_outputs[0] )
											UpperCamelCase_:							Union[str, Any]		=  [output_layers[self.config.num_hidden_layers - 1](snake_case_ )]
							else:
											UpperCamelCase_:							Dict		=  0
											UpperCamelCase_:							Optional[int]		=  None
											UpperCamelCase_:							Optional[int]		=  0
											for i in range(self.config.num_hidden_layers ):
															calculated_layer_num += 1
															UpperCamelCase_:							Optional[int]		=  self.encoder.adaptive_forward(
															    snake_case_  ,		current_layer=snake_case_  ,		attention_mask=snake_case_  ,		head_mask=snake_case_ )

															UpperCamelCase_:							Union[str, Any]		=  self.pooler(snake_case_ )
															UpperCamelCase_:							Optional[Any]		=  output_layers[i](snake_case_ )
															if regression:
																			UpperCamelCase_:							List[Any]		=  logits.detach()
																			if patient_result is not None:
																							UpperCamelCase_:							int		=  patient_result.detach()
																			if (patient_result is not None) and torch.abs(patient_result - labels ) < self.regression_threshold:
																							patient_counter += 1
																			else:
																							UpperCamelCase_:							Any		=  0
															else:
																			UpperCamelCase_:							Any		=  logits.detach().argmax(dim=1 )
																			if patient_result is not None:
																							UpperCamelCase_:							str		=  patient_result.detach().argmax(dim=1 )
																			if (patient_result is not None) and torch.all(labels.eq(snake_case_ ) ):
																							patient_counter += 1
																			else:
																							UpperCamelCase_:							Union[str, Any]		=  0

															UpperCamelCase_:							List[Any]		=  logits
															if patient_counter == self.patience:
																			break
											UpperCamelCase_:							Any		=  [patient_result]
											self.inference_layers_num += calculated_layer_num
											self.inference_instances_num += 1

							return res




@add_start_docstrings(
    """Bert Model transformer with PABEE and a sequence classification/regression head on top (a linear layer on top of
    the pooled output) e.g. for GLUE tasks. """       ,    __UpperCamelCase       ,    )
class     _UpperCamelCase			(   __UpperCamelCase					):






			'''simple docstring'''

			def __init__(    self	:    Any  ,		snake_case_	:    Union[str, Any] ):
							super().__init__(snake_case_ )
							UpperCamelCase_:							Optional[int]		=  config.num_labels

							UpperCamelCase_:							Tuple		=  BertModelWithPabee(snake_case_ )
							UpperCamelCase_:							List[str]		=  nn.Dropout(config.hidden_dropout_prob )
							UpperCamelCase_:							Any		=  nn.ModuleList(
							    [nn.Linear(config.hidden_size  ,		self.config.num_labels ) for _ in range(config.num_hidden_layers )] )

							self.init_weights()




			@add_start_docstrings_to_model_forward(snake_case_ )
			def  lowerCAmelCase__		(    self	:    Tuple  ,		snake_case_	:    int=None  ,		snake_case_	:    int=None  ,		snake_case_	:    Tuple=None  ,		snake_case_	:    Any=None  ,		snake_case_	:    Union[str, Any]=None  ,		snake_case_	:    Union[str, Any]=None  ,		snake_case_	:    Tuple=None  ,		):
							UpperCamelCase_:							Dict		=  self.bert(
							    input_ids=snake_case_  ,		attention_mask=snake_case_  ,		token_type_ids=snake_case_  ,		position_ids=snake_case_  ,		head_mask=snake_case_  ,		inputs_embeds=snake_case_  ,		output_dropout=self.dropout  ,		output_layers=self.classifiers  ,		regression=self.num_labels == 1  ,		)

							UpperCamelCase_:							int		=  (logits[-1],)

							if labels is not None:
											UpperCamelCase_:							Any		=  None
											UpperCamelCase_:							Union[str, Any]		=  0
											for ix, logits_item in enumerate(snake_case_ ):
															if self.num_labels == 1:
																			#  We are doing regression
																			UpperCamelCase_:							Dict		=  MSELoss()
																			UpperCamelCase_:							Dict		=  loss_fct(logits_item.view(-1 )  ,		labels.view(-1 ) )
															else:
																			UpperCamelCase_:							Optional[Any]		=  CrossEntropyLoss()
																			UpperCamelCase_:							int		=  loss_fct(logits_item.view(-1  ,		self.num_labels )  ,		labels.view(-1 ) )
															if total_loss is None:
																			UpperCamelCase_:							Any		=  loss
															else:
																			total_loss += loss * (ix + 1)
															total_weights += ix + 1
											UpperCamelCase_:							Tuple		=  (total_loss / total_weights,) + outputs

							return outputs








import io
import itertools
import json
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa
import pyarrow.json as paj

import datasets
from datasets.table import table_cast
from datasets.utils.file_utils import readline


lowerCamelCase_  :	Optional[Any] 				=					datasets.utils.logging.get_logger(__name__)

@dataclass
class     _UpperCamelCase			(   datasets.BuilderConfig					):






			'''simple docstring'''

			__UpperCamelCase :  Optional[datasets.Features]    =	None
			__UpperCamelCase :  str    =	"utf-8"
			__UpperCamelCase :  Optional[str]    =	None
			__UpperCamelCase :  Optional[str]    =	None
			__UpperCamelCase :  bool    =	True  # deprecated
			__UpperCamelCase :  Optional[int]    =	None  # deprecated
			__UpperCamelCase :  int    =	10 << 20  # 10MB
			__UpperCamelCase :  Optional[bool]    =	None




class     _UpperCamelCase			(   datasets.ArrowBasedBuilder					):






			'''simple docstring'''

			__UpperCamelCase :  Tuple    =	JsonConfig

			def  lowerCAmelCase__		(    self	:    int ):
							if self.config.block_size is not None:
											logger.warning("""The JSON loader parameter `block_size` is deprecated. Please use `chunksize` instead""" )
											UpperCamelCase_:							List[str]		=  self.config.block_size
							if self.config.use_threads is not True:
											logger.warning(
											    """The JSON loader parameter `use_threads` is deprecated and doesn't have any effect anymore.""" )
							if self.config.newlines_in_values is not None:
											raise ValueError("""The JSON loader parameter `newlines_in_values` is no longer supported""" )
							return datasets.DatasetInfo(features=self.config.features )

			def  lowerCAmelCase__		(    self	:    Dict  ,		snake_case_	:    str ):
							if not self.config.data_files:
											raise ValueError(f'''At least one data file must be specified, but got data_files={self.config.data_files}''' )
							UpperCamelCase_:							Dict		=  dl_manager.download_and_extract(self.config.data_files )
							if isinstance(snake_case_  ,		(str, list, tuple) ):
											UpperCamelCase_:							List[Any]		=  data_files
											if isinstance(snake_case_  ,		snake_case_ ):
															UpperCamelCase_:							str		=  [files]
											UpperCamelCase_:							Any		=  [dl_manager.iter_files(snake_case_ ) for file in files]
											return [datasets.SplitGenerator(name=datasets.Split.TRAIN  ,		gen_kwargs={"""files""": files} )]
							UpperCamelCase_:							Dict		=  []
							for split_name, files in data_files.items():
											if isinstance(snake_case_  ,		snake_case_ ):
															UpperCamelCase_:							Tuple		=  [files]
											UpperCamelCase_:							Optional[int]		=  [dl_manager.iter_files(snake_case_ ) for file in files]
											splits.append(datasets.SplitGenerator(name=snake_case_  ,		gen_kwargs={"""files""": files} ) )
							return splits

			def  lowerCAmelCase__		(    self	:    str  ,		snake_case_	:    pa.Table ):
							if self.config.features is not None:
											# adding missing columns
											for column_name in set(self.config.features ) - set(pa_table.column_names ):
															UpperCamelCase_:							Union[str, Any]		=  self.config.features.arrow_schema.field(snake_case_ ).type
															UpperCamelCase_:							Tuple		=  pa_table.append_column(snake_case_  ,		pa.array([None] * len(snake_case_ )  ,		type=snake_case_ ) )
											# more expensive cast to support nested structures with keys in a different order
											# allows str <-> int/float or str to Audio for example
											UpperCamelCase_:							int		=  table_cast(snake_case_  ,		self.config.features.arrow_schema )
							return pa_table




			def  lowerCAmelCase__		(    self	:    Dict  ,		snake_case_	:    Optional[Any] ):
							for file_idx, file in enumerate(itertools.chain.from_iterable(snake_case_ ) ):
											# If the file is one json object and if we need to look at the list of items in one specific field
											if self.config.field is not None:
															with open(snake_case_  ,		encoding=self.config.encoding  ,		errors=self.config.encoding_errors ) as f:
																			UpperCamelCase_:							Dict		=  json.load(snake_case_ )

															# We keep only the field we are interested in
															UpperCamelCase_:							Optional[int]		=  dataset[self.config.field]

															# We accept two format: a list of dicts or a dict of lists
															if isinstance(snake_case_  ,		(list, tuple) ):
																			UpperCamelCase_:							Optional[int]		=  set().union(*[row.keys() for row in dataset] )
																			UpperCamelCase_:							int		=  {col: [row.get(snake_case_ ) for row in dataset] for col in keys}
															else:
																			UpperCamelCase_:							Optional[int]		=  dataset
															UpperCamelCase_:							List[str]		=  pa.Table.from_pydict(snake_case_ )
															yield file_idx, self._cast_table(snake_case_ )

											# If the file has one json object per line
											else:
															with open(snake_case_  ,		"""rb""" ) as f:
																			UpperCamelCase_:							Optional[int]		=  0
																			# Use block_size equal to the chunk size divided by 32 to leverage multithreading
																			# Set a default minimum value of 16kB if the chunk size is really small
																			UpperCamelCase_:							Optional[int]		=  max(self.config.chunksize // 32  ,		16 << 10 )
																			UpperCamelCase_:							Tuple		=  (
																			    self.config.encoding_errors if self.config.encoding_errors is not None else """strict"""
																			)
																			while True:
																							UpperCamelCase_:							int		=  f.read(self.config.chunksize )
																							if not batch:
																											break
																							# Finish current line
																							try:
																											batch += f.readline()
																							except (AttributeError, io.UnsupportedOperation):
																											batch += readline(snake_case_ )
																							# PyArrow only accepts utf-8 encoded bytes
																							if self.config.encoding != "utf-8":
																											UpperCamelCase_:							Tuple		=  batch.decode(self.config.encoding  ,		errors=snake_case_ ).encode("""utf-8""" )
																							try:
																											while True:
																															try:
																																			UpperCamelCase_:							Tuple		=  paj.read_json(
																																			    io.BytesIO(snake_case_ )  ,		read_options=paj.ReadOptions(block_size=snake_case_ ) )
																																			break
																															except (pa.ArrowInvalid, pa.ArrowNotImplementedError) as e:
																																			if (
																																			    isinstance(snake_case_  ,		pa.ArrowInvalid )
																																			    and "straddling" not in str(snake_case_ )
																																			    or block_size > len(snake_case_ )
																																			):
																																							raise
																																			else:
																																							# Increase the block size in case it was too small.
																																							# The block size will be reset for the next file.
																																							logger.debug(
																																							    f'''Batch of {len(snake_case_ )} bytes couldn\'t be parsed with block_size={block_size}. Retrying with block_size={block_size * 2}.''' )
																																							block_size *= 2
																							except pa.ArrowInvalid as e:
																											try:
																															with open(
																															    snake_case_  ,		encoding=self.config.encoding  ,		errors=self.config.encoding_errors ) as f:
																																			UpperCamelCase_:							Optional[Any]		=  json.load(snake_case_ )
																											except json.JSONDecodeError:
																															logger.error(f'''Failed to read file \'{file}\' with error {type(snake_case_ )}: {e}''' )
																															raise e
																											# If possible, parse the file as a list of json objects and exit the loop
																											if isinstance(snake_case_  ,		snake_case_ ):  # list is the only sequence type supported in JSON
																															try:
																																			UpperCamelCase_:							Any		=  set().union(*[row.keys() for row in dataset] )
																																			UpperCamelCase_:							List[str]		=  {col: [row.get(snake_case_ ) for row in dataset] for col in keys}
																																			UpperCamelCase_:							int		=  pa.Table.from_pydict(snake_case_ )
																															except (pa.ArrowInvalid, AttributeError) as e:
																																			logger.error(f'''Failed to read file \'{file}\' with error {type(snake_case_ )}: {e}''' )
																																			raise ValueError(f'''Not able to read records in the JSON file at {file}.''' ) from None
																															yield file_idx, self._cast_table(snake_case_ )
																															break
																											else:
																															logger.error(f'''Failed to read file \'{file}\' with error {type(snake_case_ )}: {e}''' )
																															raise ValueError(
																															    f'''Not able to read records in the JSON file at {file}. '''
																															    f'''You should probably indicate the field of the JSON file containing your records. '''
																															    f'''This JSON file contain the following fields: {str(list(dataset.keys() ) )}. '''
																															    f'''Select the correct one and provide it as `field=\'XXX\'` to the dataset loading method. ''' ) from None
                        # Uncomment for debugging (will print the Arrow table size and elements)
                        # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
                        # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
																							yield (file_idx, batch_idx), self._cast_table(snake_case_ )
																							batch_idx += 1






