
from typing import Optional, Tuple, Union

import tensorflow as tf

from ...activations_tf import ACTaFN
from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
from ...modeling_tf_outputs import (
    TFBaseModelOutputWithNoAttention,
    TFBaseModelOutputWithPoolingAndNoAttention,
    TFSequenceClassifierOutput,
)
from ...modeling_tf_utils import TFPreTrainedModel, TFSequenceClassificationLoss, keras_serializable, unpack_inputs
from ...tf_utils import shape_list
from ...utils import logging
from .configuration_regnet import RegNetConfig


_UpperCAmelCase      :				str          = logging.get_logger(__name__)

# General docstring
_UpperCAmelCase      :				Any          = "RegNetConfig"

# Base docstring
_UpperCAmelCase      :				str          = "facebook/regnet-y-040"
_UpperCAmelCase      :				Optional[int]          = [1, 1_088, 7, 7]

# Image classification docstring
_UpperCAmelCase      :				Optional[int]          = "facebook/regnet-y-040"
_UpperCAmelCase      :				str          = "tabby, tabby cat"

_UpperCAmelCase      :				List[Any]          = [
    "facebook/regnet-y-040",
    # See all regnet models at https://huggingface.co/models?filter=regnet
]


class        lowercase      (	tf.keras.layers.Layer     ):



		def __init__(				self       ,   A_       ,   A_ = 3       ,   A_ = 1       ,   A_ = 1       ,   A_ = "relu"       ,   **A_       ,   )						->      Optional[Any]:
				"""simple docstring"""
				super().__init__(**A_ )
				# The padding and conv has been verified in
				# https://colab.research.google.com/gist/sayakpaul/854bc10eeaf21c9ee2119e0b9f3841a7/scratchpad.ipynb
				UpperCamelCase						=  tf.keras.layers.ZeroPaddingaD(padding=kernel_size // 2 )
				UpperCamelCase						=  tf.keras.layers.ConvaD(
				    filters=A_       ,   kernel_size=A_       ,   strides=A_       ,   padding='VALID'       ,   groups=A_       ,   use_bias=A_       ,   name='convolution'       ,   )
				UpperCamelCase						=  tf.keras.layers.BatchNormalization(epsilon=1e-5       ,   momentum=0.9       ,   name='normalization' )
				UpperCamelCase						=  ACTaFN[activation] if activation is not None else tf.identity



		def __UpperCamelCase				(				self       ,   A_ )						->      str:
				"""simple docstring"""
				UpperCamelCase						=  self.convolution(self.padding(A_ ) )
				UpperCamelCase						=  self.normalization(A_ )
				UpperCamelCase						=  self.activation(A_ )
				return hidden_state






class        lowercase      (	tf.keras.layers.Layer     ):



		def __init__(				self       ,   A_       ,   **A_ )						->      Tuple:
				"""simple docstring"""
				super().__init__(**A_ )
				UpperCamelCase						=  config.num_channels
				UpperCamelCase						=  TFRegNetConvLayer(
				    out_channels=config.embedding_size       ,   kernel_size=3       ,   stride=2       ,   activation=config.hidden_act       ,   name='embedder'       ,   )



		def __UpperCamelCase				(				self       ,   A_ )						->      str:
				"""simple docstring"""
				UpperCamelCase						=  shape_list(A_ )[1]
				if tf.executing_eagerly() and num_channels != self.num_channels:
						raise ValueError(
						    'Make sure that the channel dimension of the pixel values match with the one set in the configuration.' )

				# When running on CPU, `tf.keras.layers.Conv2D` doesn't support `NCHW` format.
				# So change the input format from `NCHW` to `NHWC`.
				# shape = (batch_size, in_height, in_width, in_channels=num_channels)
				UpperCamelCase						=  tf.transpose(A_       ,   perm=(0, 2, 3, 1) )
				UpperCamelCase						=  self.embedder(A_ )
				return hidden_state






class        lowercase      (	tf.keras.layers.Layer     ):



		def __init__(				self       ,   A_       ,   A_ = 2       ,   **A_ )						->      Dict:
				"""simple docstring"""
				super().__init__(**A_ )
				UpperCamelCase						=  tf.keras.layers.ConvaD(
				    filters=A_       ,   kernel_size=1       ,   strides=A_       ,   use_bias=A_       ,   name='convolution' )
				UpperCamelCase						=  tf.keras.layers.BatchNormalization(epsilon=1e-5       ,   momentum=0.9       ,   name='normalization' )



		def __UpperCamelCase				(				self       ,   A_       ,   A_ = False )						->      tf.Tensor:
				"""simple docstring"""
				return self.normalization(self.convolution(A_ )       ,   training=A_ )






class        lowercase      (	tf.keras.layers.Layer     ):



		def __init__(				self       ,   A_       ,   A_       ,   **A_ )						->      Any:
				"""simple docstring"""
				super().__init__(**A_ )
				UpperCamelCase						=  tf.keras.layers.GlobalAveragePoolingaD(keepdims=A_       ,   name='pooler' )
				UpperCamelCase						=  [
				    tf.keras.layers.ConvaD(filters=A_       ,   kernel_size=1       ,   activation='relu'       ,   name='attention.0' ),
				    tf.keras.layers.ConvaD(filters=A_       ,   kernel_size=1       ,   activation='sigmoid'       ,   name='attention.2' ),
				]



		def __UpperCamelCase				(				self       ,   A_ )						->      List[str]:
				"""simple docstring"""
				# [batch_size, h, w, num_channels] -> [batch_size, 1, 1, num_channels]
				UpperCamelCase						=  self.pooler(A_ )
				for layer_module in self.attention:
						UpperCamelCase						=  layer_module(A_ )
				UpperCamelCase						=  hidden_state * pooled
				return hidden_state






class        lowercase      (	tf.keras.layers.Layer     ):



		def __init__(				self       ,   A_       ,   A_       ,   A_       ,   A_ = 1       ,   **A_ )						->      Dict:
				"""simple docstring"""
				super().__init__(**A_ )
				UpperCamelCase						=  in_channels != out_channels or stride != 1
				UpperCamelCase						=  max(1       ,   out_channels // config.groups_width )
				UpperCamelCase						=  (
				    TFRegNetShortCut(A_       ,   stride=A_       ,   name='shortcut' )
				    if should_apply_shortcut
				    else tf.keras.layers.Activation('linear'       ,   name='shortcut' )
				)
				# `self.layers` instead of `self.layer` because that is a reserved argument.
				UpperCamelCase						=  [
				    TFRegNetConvLayer(A_       ,   kernel_size=1       ,   activation=config.hidden_act       ,   name='layer.0' ),
				    TFRegNetConvLayer(
				        A_       ,   stride=A_       ,   groups=A_       ,   activation=config.hidden_act       ,   name='layer.1' ),
				    TFRegNetConvLayer(A_       ,   kernel_size=1       ,   activation=A_       ,   name='layer.2' ),
				]
				UpperCamelCase						=  ACTaFN[config.hidden_act]



		def __UpperCamelCase				(				self       ,   A_ )						->      str:
				"""simple docstring"""
				UpperCamelCase						=  hidden_state
				for layer_module in self.layers:
						UpperCamelCase						=  layer_module(A_ )
				UpperCamelCase						=  self.shortcut(A_ )
				hidden_state += residual
				UpperCamelCase						=  self.activation(A_ )
				return hidden_state






class        lowercase      (	tf.keras.layers.Layer     ):



		def __init__(				self       ,   A_       ,   A_       ,   A_       ,   A_ = 1       ,   **A_ )						->      Optional[int]:
				"""simple docstring"""
				super().__init__(**A_ )
				UpperCamelCase						=  in_channels != out_channels or stride != 1
				UpperCamelCase						=  max(1       ,   out_channels // config.groups_width )
				UpperCamelCase						=  (
				    TFRegNetShortCut(A_       ,   stride=A_       ,   name='shortcut' )
				    if should_apply_shortcut
				    else tf.keras.layers.Activation('linear'       ,   name='shortcut' )
				)
				UpperCamelCase						=  [
				    TFRegNetConvLayer(A_       ,   kernel_size=1       ,   activation=config.hidden_act       ,   name='layer.0' ),
				    TFRegNetConvLayer(
				        A_       ,   stride=A_       ,   groups=A_       ,   activation=config.hidden_act       ,   name='layer.1' ),
				    TFRegNetSELayer(A_       ,   reduced_channels=int(round(in_channels / 4 ) )       ,   name='layer.2' ),
				    TFRegNetConvLayer(A_       ,   kernel_size=1       ,   activation=A_       ,   name='layer.3' ),
				]
				UpperCamelCase						=  ACTaFN[config.hidden_act]



		def __UpperCamelCase				(				self       ,   A_ )						->      Optional[Any]:
				"""simple docstring"""
				UpperCamelCase						=  hidden_state
				for layer_module in self.layers:
						UpperCamelCase						=  layer_module(A_ )
				UpperCamelCase						=  self.shortcut(A_ )
				hidden_state += residual
				UpperCamelCase						=  self.activation(A_ )
				return hidden_state






class        lowercase      (	tf.keras.layers.Layer     ):



		def __init__(				self       ,   A_       ,   A_       ,   A_       ,   A_ = 2       ,   A_ = 2       ,   **A_ )						->      Any:
				"""simple docstring"""
				super().__init__(**A_ )

				UpperCamelCase						=  TFRegNetXLayer if config.layer_type == 'x' else TFRegNetYLayer
				UpperCamelCase						=  [
				    # downsampling is done in the first layer with stride of 2
				    layer(A_       ,   A_       ,   A_       ,   stride=A_       ,   name='layers.0' ),
				    *[layer(A_       ,   A_       ,   A_       ,   name=F'''layers.{i+1}''' ) for i in range(depth - 1 )],
				]



		def __UpperCamelCase				(				self       ,   A_ )						->      int:
				"""simple docstring"""
				for layer_module in self.layers:
						UpperCamelCase						=  layer_module(A_ )
				return hidden_state






class        lowercase      (	tf.keras.layers.Layer     ):



		def __init__(				self       ,   A_       ,   **A_ )						->      Dict:
				"""simple docstring"""
				super().__init__(**A_ )
				UpperCamelCase						=  []
				# based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
				self.stages.append(
				    TFRegNetStage(
				        A_       ,   config.embedding_size       ,   config.hidden_sizes[0]       ,   stride=2 if config.downsample_in_first_stage else 1       ,   depth=config.depths[0]       ,   name='stages.0'       ,   ) )
				UpperCamelCase						=  zip(config.hidden_sizes       ,   config.hidden_sizes[1:] )
				for i, ((in_channels, out_channels), depth) in enumerate(zip(A_       ,   config.depths[1:] ) ):
						self.stages.append(TFRegNetStage(A_       ,   A_       ,   A_       ,   depth=A_       ,   name=F'''stages.{i+1}''' ) )



		def __UpperCamelCase				(				self       ,   A_       ,   A_ = False       ,   A_ = True )						->      TFBaseModelOutputWithNoAttention:
				"""simple docstring"""
				UpperCamelCase						=  () if output_hidden_states else None

				for stage_module in self.stages:
						if output_hidden_states:
								UpperCamelCase						=  hidden_states + (hidden_state,)

						UpperCamelCase						=  stage_module(A_ )

				if output_hidden_states:
						UpperCamelCase						=  hidden_states + (hidden_state,)

				if not return_dict:
						return tuple(v for v in [hidden_state, hidden_states] if v is not None )

				return TFBaseModelOutputWithNoAttention(last_hidden_state=A_       ,   hidden_states=A_ )






@keras_serializable
class        lowercase      (	tf.keras.layers.Layer     ):
		__lowercase : Union[str, Any]			  =       RegNetConfig



		def __init__(				self       ,   A_       ,   **A_ )						->      Tuple:
				"""simple docstring"""
				super().__init__(**A_ )
				UpperCamelCase						=  config
				UpperCamelCase						=  TFRegNetEmbeddings(A_       ,   name='embedder' )
				UpperCamelCase						=  TFRegNetEncoder(A_       ,   name='encoder' )
				UpperCamelCase						=  tf.keras.layers.GlobalAveragePoolingaD(keepdims=A_       ,   name='pooler' )



		@unpack_inputs
		def __UpperCamelCase				(				self       ,   A_       ,   A_ = None       ,   A_ = None       ,   A_ = False       ,   )						->      TFBaseModelOutputWithPoolingAndNoAttention:
				"""simple docstring"""
				UpperCamelCase						=  (
				    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
				)
				UpperCamelCase						=  return_dict if return_dict is not None else self.config.use_return_dict

				UpperCamelCase						=  self.embedder(A_       ,   training=A_ )

				UpperCamelCase						=  self.encoder(
				    A_       ,   output_hidden_states=A_       ,   return_dict=A_       ,   training=A_ )

				UpperCamelCase						=  encoder_outputs[0]
				UpperCamelCase						=  self.pooler(A_ )

				# Change to NCHW output format have uniformity in the modules
				UpperCamelCase						=  tf.transpose(A_       ,   perm=(0, 3, 1, 2) )
				UpperCamelCase						=  tf.transpose(A_       ,   perm=(0, 3, 1, 2) )

				# Change the other hidden state outputs to NCHW as well
				if output_hidden_states:
						UpperCamelCase						=  tuple([tf.transpose(A_       ,   perm=(0, 3, 1, 2) ) for h in encoder_outputs[1]] )

				if not return_dict:
						return (last_hidden_state, pooled_output) + encoder_outputs[1:]

				return TFBaseModelOutputWithPoolingAndNoAttention(
				    last_hidden_state=A_       ,   pooler_output=A_       ,   hidden_states=hidden_states if output_hidden_states else encoder_outputs.hidden_states       ,   )






class        lowercase      (	_SCREAMING_SNAKE_CASE     ):
		__lowercase : str			  =       RegNetConfig
		__lowercase : List[Any]			  =       "regnet"
		__lowercase : List[Any]			  =       "pixel_values"



		@property
		def __UpperCamelCase				(				self )						->      int:
				"""simple docstring"""
				return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 224, 224)       ,   dtype=tf.floataa )}






_UpperCAmelCase      :				str          = R"\n    Parameters:\n    This model is a Tensorflow\n    [tf.keras.layers.Layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer) sub-class. Use it as a\n    regular Tensorflow Module and refer to the Tensorflow documentation for all matter related to general usage and\n    behavior.\n        config ([`RegNetConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.\n"

_UpperCAmelCase      :				Dict          = R"\n    Args:\n        pixel_values (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`ConveNextImageProcessor.__call__`] for details.\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for\n            more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n"


@add_start_docstrings(
    "The bare RegNet model outputting raw features without any specific head on top."						,     _SCREAMING_SNAKE_CASE						,     )
class        lowercase      (	_SCREAMING_SNAKE_CASE     ):



		def __init__(				self       ,   A_       ,   *A_       ,   **A_ )						->      Any:
				"""simple docstring"""
				super().__init__(A_       ,   *A_       ,   **A_ )
				UpperCamelCase						=  TFRegNetMainLayer(A_       ,   name='regnet' )



		@unpack_inputs
		@add_start_docstrings_to_model_forward(A_ )
		@add_code_sample_docstrings(
		    checkpoint=_CHECKPOINT_FOR_DOC       ,   output_type=A_       ,   config_class=_CONFIG_FOR_DOC       ,   modality='vision'       ,   expected_output=_EXPECTED_OUTPUT_SHAPE       ,   )
		def __UpperCamelCase				(				self       ,   A_       ,   A_ = None       ,   A_ = None       ,   A_=False       ,   )						->      Union[TFBaseModelOutputWithPoolingAndNoAttention, Tuple[tf.Tensor]]:
				"""simple docstring"""
				UpperCamelCase						=  (
				    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
				)
				UpperCamelCase						=  return_dict if return_dict is not None else self.config.use_return_dict

				UpperCamelCase						=  self.regnet(
				    pixel_values=A_       ,   output_hidden_states=A_       ,   return_dict=A_       ,   training=A_       ,   )
				if not return_dict:
						return (outputs[0],) + outputs[1:]

				return TFBaseModelOutputWithPoolingAndNoAttention(
				    last_hidden_state=outputs.last_hidden_state       ,   pooler_output=outputs.pooler_output       ,   hidden_states=outputs.hidden_states       ,   )






@add_start_docstrings(
    "\n    RegNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    "						,     _SCREAMING_SNAKE_CASE						,     )
class        lowercase      (	_SCREAMING_SNAKE_CASE						,     _SCREAMING_SNAKE_CASE     ):



		def __init__(				self       ,   A_       ,   *A_       ,   **A_ )						->      Optional[Any]:
				"""simple docstring"""
				super().__init__(A_       ,   *A_       ,   **A_ )
				UpperCamelCase						=  config.num_labels
				UpperCamelCase						=  TFRegNetMainLayer(A_       ,   name='regnet' )
				# classification head
				UpperCamelCase						=  [
				    tf.keras.layers.Flatten(),
				    tf.keras.layers.Dense(config.num_labels       ,   name='classifier.1' ) if config.num_labels > 0 else tf.identity,
				]



		@unpack_inputs
		@add_start_docstrings_to_model_forward(A_ )
		@add_code_sample_docstrings(
		    checkpoint=_IMAGE_CLASS_CHECKPOINT       ,   output_type=A_       ,   config_class=_CONFIG_FOR_DOC       ,   expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT       ,   )
		def __UpperCamelCase				(				self       ,   A_ = None       ,   A_ = None       ,   A_ = None       ,   A_ = None       ,   A_=False       ,   )						->      Union[TFSequenceClassifierOutput, Tuple[tf.Tensor]]:
				"""simple docstring"""
				UpperCamelCase						=  (
				    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
				)
				UpperCamelCase						=  return_dict if return_dict is not None else self.config.use_return_dict

				UpperCamelCase						=  self.regnet(
				    A_       ,   output_hidden_states=A_       ,   return_dict=A_       ,   training=A_ )

				UpperCamelCase						=  outputs.pooler_output if return_dict else outputs[1]

				UpperCamelCase						=  self.classifier[0](A_ )
				UpperCamelCase						=  self.classifier[1](A_ )

				UpperCamelCase						=  None if labels is None else self.hf_compute_loss(labels=A_       ,   logits=A_ )

				if not return_dict:
						UpperCamelCase						=  (logits,) + outputs[2:]
						return ((loss,) + output) if loss is not None else output

				return TFSequenceClassifierOutput(loss=A_       ,   logits=A_       ,   hidden_states=outputs.hidden_states )


from abc import ABC, abstractmethod
from typing import List, Optional


class        lowercase      (	_SCREAMING_SNAKE_CASE     ):



		def __init__(				self )						->      Optional[Any]:
				"""simple docstring"""
				# test for the above condition
				self.test()



		def __UpperCamelCase				(				self )						->      Dict:
				"""simple docstring"""
				UpperCamelCase						=  0
				UpperCamelCase						=  False
				while not completed:
						if counter == 1:
								self.reset()
						UpperCamelCase						=  self.advance()
						if not self.does_advance(A_ ):
								raise Exception(
								    'Custom Constraint is not defined correctly. self.does_advance(self.advance()) must be true.' )

						UpperCamelCase			,						UpperCamelCase			,						UpperCamelCase						=  self.update(A_ )
						counter += 1

						if counter > 10_000:
								raise Exception('update() does not fulfill the constraint.' )

				if self.remaining() != 0:
						raise Exception('Custom Constraint is not defined correctly.' )



		@abstractmethod
		def __UpperCamelCase				(				self )						->      Optional[Any]:
				"""simple docstring"""
				raise NotImplementedError(
				    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.''' )



		@abstractmethod
		def __UpperCamelCase				(				self       ,   A_ )						->      str:
				"""simple docstring"""
				raise NotImplementedError(
				    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.''' )



		@abstractmethod
		def __UpperCamelCase				(				self       ,   A_ )						->      int:
				"""simple docstring"""
				raise NotImplementedError(
				    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.''' )



		@abstractmethod
		def __UpperCamelCase				(				self )						->      Any:
				"""simple docstring"""
				raise NotImplementedError(
				    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.''' )



		@abstractmethod
		def __UpperCamelCase				(				self )						->      str:
				"""simple docstring"""
				raise NotImplementedError(
				    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.''' )



		@abstractmethod
		def __UpperCamelCase				(				self       ,   A_=False )						->      int:
				"""simple docstring"""
				raise NotImplementedError(
				    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.''' )






class        lowercase      (	_SCREAMING_SNAKE_CASE     ):



		def __init__(				self       ,   A_ )						->      Any:
				"""simple docstring"""
				super(A_       ,   self ).__init__()

				if not isinstance(A_       ,   A_ ) or len(A_ ) == 0:
						raise ValueError(F'''`token_ids` has to be a non-empty list, but is {token_ids}.''' )
				if any((not isinstance(A_       ,   A_ ) or token_id < 0) for token_id in token_ids ):
						raise ValueError(F'''Each list in `token_ids` has to be a list of positive integers, but is {token_ids}.''' )

				UpperCamelCase						=  token_ids

				UpperCamelCase						=  len(self.token_ids )
				UpperCamelCase						=  -1  # the index of the currently fulfilled step
				UpperCamelCase						=  False



		def __UpperCamelCase				(				self )						->      Optional[Any]:
				"""simple docstring"""
				if self.completed:
						return None
				return self.token_ids[self.fulfilled_idx + 1]



		def __UpperCamelCase				(				self       ,   A_ )						->      Optional[int]:
				"""simple docstring"""
				if not isinstance(A_       ,   A_ ):
						raise ValueError(F'''`token_id` has to be an `int`, but is {token_id} of type {type(A_ )}''' )

				if self.completed:
						return False

				return token_id == self.token_ids[self.fulfilled_idx + 1]



		def __UpperCamelCase				(				self       ,   A_ )						->      Optional[int]:
				"""simple docstring"""
				if not isinstance(A_       ,   A_ ):
						raise ValueError(F'''`token_id` has to be an `int`, but is {token_id} of type {type(A_ )}''' )

				UpperCamelCase						=  False
				UpperCamelCase						=  False
				UpperCamelCase						=  False

				if self.does_advance(A_ ):
						self.fulfilled_idx += 1
						UpperCamelCase						=  True
						if self.fulfilled_idx == (self.seqlen - 1):
								UpperCamelCase						=  True
						UpperCamelCase						=  completed
				else:
						# failed to make progress.
						UpperCamelCase						=  True
						self.reset()
				return stepped, completed, reset



		def __UpperCamelCase				(				self )						->      Dict:
				"""simple docstring"""
				UpperCamelCase						=  False
				UpperCamelCase						=  0



		def __UpperCamelCase				(				self )						->      int:
				"""simple docstring"""
				return self.seqlen - (self.fulfilled_idx + 1)



		def __UpperCamelCase				(				self       ,   A_=False )						->      Union[str, Any]:
				"""simple docstring"""
				UpperCamelCase						=  PhrasalConstraint(self.token_ids )

				if stateful:
						UpperCamelCase						=  self.seqlen
						UpperCamelCase						=  self.fulfilled_idx
						UpperCamelCase						=  self.completed

				return new_constraint






class        lowercase      :



		def __init__(				self       ,   A_       ,   A_=True )						->      List[Any]:
				"""simple docstring"""
				UpperCamelCase						=  max([len(A_ ) for one in nested_token_ids] )

				UpperCamelCase						=  {}
				for token_ids in nested_token_ids:
						UpperCamelCase						=  root
						for tidx, token_id in enumerate(A_ ):
								if token_id not in level:
										UpperCamelCase						=  {}

								UpperCamelCase						=  level[token_id]

				if no_subsets and self.has_subsets(A_       ,   A_ ):
						raise ValueError(
						    'Each list in `nested_token_ids` can\'t be a complete subset of another list, but is'
						    F''' {nested_token_ids}.''' )

				UpperCamelCase						=  root



		def __UpperCamelCase				(				self       ,   A_ )						->      Optional[Any]:
				"""simple docstring"""
				UpperCamelCase						=  self.trie

				for current_token in current_seq:
						UpperCamelCase						=  start[current_token]

				UpperCamelCase						=  list(start.keys() )

				return next_tokens



		def __UpperCamelCase				(				self       ,   A_ )						->      Union[str, Any]:
				"""simple docstring"""
				UpperCamelCase						=  self.next_tokens(A_ )

				return len(A_ ) == 0



		def __UpperCamelCase				(				self       ,   A_ )						->      List[str]:
				"""simple docstring"""
				UpperCamelCase						=  list(root.values() )
				if len(A_ ) == 0:
						return 1
				else:
						return sum([self.count_leaves(A_ ) for nn in next_nodes] )



		def __UpperCamelCase				(				self       ,   A_       ,   A_ )						->      Optional[int]:
				"""simple docstring"""
				UpperCamelCase						=  self.count_leaves(A_ )
				return len(A_ ) != leaf_count






class        lowercase      (	_SCREAMING_SNAKE_CASE     ):



		def __init__(				self       ,   A_ )						->      str:
				"""simple docstring"""
				super(A_       ,   self ).__init__()

				if not isinstance(A_       ,   A_ ) or len(A_ ) == 0:
						raise ValueError(F'''`nested_token_ids` has to be a non-empty list, but is {nested_token_ids}.''' )
				if any(not isinstance(A_       ,   A_ ) for token_ids in nested_token_ids ):
						raise ValueError(F'''`nested_token_ids` has to be a list of lists, but is {nested_token_ids}.''' )
				if any(
				    any((not isinstance(A_       ,   A_ ) or token_id < 0) for token_id in token_ids )
				    for token_ids in nested_token_ids ):
						raise ValueError(
						    F'''Each list in `nested_token_ids` has to be a list of positive integers, but is {nested_token_ids}.''' )

				UpperCamelCase						=  DisjunctiveTrie(A_ )
				UpperCamelCase						=  nested_token_ids

				UpperCamelCase						=  self.trie.max_height
				UpperCamelCase						=  []
				UpperCamelCase						=  False



		def __UpperCamelCase				(				self )						->      Union[str, Any]:
				"""simple docstring"""
				UpperCamelCase						=  self.trie.next_tokens(self.current_seq )

				if len(A_ ) == 0:
						return None
				else:
						return token_list



		def __UpperCamelCase				(				self       ,   A_ )						->      Optional[Any]:
				"""simple docstring"""
				if not isinstance(A_       ,   A_ ):
						raise ValueError(F'''`token_id` is supposed to be type `int`, but is {token_id} of type {type(A_ )}''' )

				UpperCamelCase						=  self.trie.next_tokens(self.current_seq )

				return token_id in next_tokens



		def __UpperCamelCase				(				self       ,   A_ )						->      Optional[Any]:
				"""simple docstring"""
				if not isinstance(A_       ,   A_ ):
						raise ValueError(F'''`token_id` is supposed to be type `int`, but is {token_id} of type {type(A_ )}''' )

				UpperCamelCase						=  False
				UpperCamelCase						=  False
				UpperCamelCase						=  False

				if self.does_advance(A_ ):
						self.current_seq.append(A_ )
						UpperCamelCase						=  True
				else:
						UpperCamelCase						=  True
						self.reset()

				UpperCamelCase						=  self.trie.reached_leaf(self.current_seq )
				UpperCamelCase						=  completed

				return stepped, completed, reset



		def __UpperCamelCase				(				self )						->      str:
				"""simple docstring"""
				UpperCamelCase						=  False
				UpperCamelCase						=  []



		def __UpperCamelCase				(				self )						->      Optional[int]:
				"""simple docstring"""
				if self.completed:
						# since this can be completed without reaching max height
						return 0
				else:
						return self.seqlen - len(self.current_seq )



		def __UpperCamelCase				(				self       ,   A_=False )						->      int:
				"""simple docstring"""
				UpperCamelCase						=  DisjunctiveConstraint(self.token_ids )

				if stateful:
						UpperCamelCase						=  self.seqlen
						UpperCamelCase						=  self.current_seq
						UpperCamelCase						=  self.completed

				return new_constraint






class        lowercase      :



		def __init__(				self       ,   A_ )						->      Tuple:
				"""simple docstring"""
				UpperCamelCase						=  constraints

				# max # of steps required to fulfill a given constraint
				UpperCamelCase						=  max([c.seqlen for c in constraints] )
				UpperCamelCase						=  len(A_ )
				UpperCamelCase						=  False

				self.init_state()



		def __UpperCamelCase				(				self )						->      List[str]:
				"""simple docstring"""
				UpperCamelCase						=  []
				UpperCamelCase						=  None
				UpperCamelCase						=  [constraint.copy(stateful=A_ ) for constraint in self.constraints]



		def __UpperCamelCase				(				self )						->      Optional[int]:
				"""simple docstring"""
				UpperCamelCase						=  0
				if self.inprogress_constraint:
						# extra points for having a constraint mid-fulfilled
						add += self.max_seqlen - self.inprogress_constraint.remaining()

				return (len(self.complete_constraints ) * self.max_seqlen) + add



		def __UpperCamelCase				(				self )						->      Optional[int]:
				"""simple docstring"""
				UpperCamelCase						=  []
				if self.inprogress_constraint is None:
						for constraint in self.pending_constraints:  # "pending" == "unfulfilled yet"
								UpperCamelCase						=  constraint.advance()
								if isinstance(A_       ,   A_ ):
										token_list.append(A_ )
								elif isinstance(A_       ,   A_ ):
										token_list.extend(A_ )
				else:
						UpperCamelCase						=  self.inprogress_constraint.advance()
						if isinstance(A_       ,   A_ ):
								token_list.append(A_ )
						elif isinstance(A_       ,   A_ ):
								token_list.extend(A_ )

				if len(A_ ) == 0:
						return None
				else:
						return token_list



		def __UpperCamelCase				(				self       ,   A_ )						->      Any:
				"""simple docstring"""
				self.init_state()

				if token_ids is not None:
						for token in token_ids:
								# completes or steps **one** constraint
								UpperCamelCase			,						UpperCamelCase						=  self.add(A_ )

								# the entire list of constraints are fulfilled
								if self.completed:
										break



		def __UpperCamelCase				(				self       ,   A_ )						->      int:
				"""simple docstring"""
				if not isinstance(A_       ,   A_ ):
						raise ValueError(F'''`token_id` should be an `int`, but is `{token_id}`.''' )

				UpperCamelCase			,						UpperCamelCase						=  False, False

				if self.completed:
						UpperCamelCase						=  True
						UpperCamelCase						=  False
						return complete, stepped

				if self.inprogress_constraint is not None:
						# In the middle of fulfilling a constraint. If the `token_id` *does* makes an incremental progress to current
						# job, simply update the state

						UpperCamelCase			,						UpperCamelCase			,						UpperCamelCase						=  self.inprogress_constraint.update(A_ )
						if reset:
								# 1. If the next token breaks the progress, then we must restart.
								#     e.g. constraint = "I love pies" and sequence so far is "I love" but `token_id` == "books".

								#     But that doesn't mean we self.init_state(), since we only reset the state for this particular
								#     constraint, not the full list of constraints.

								self.pending_constraints.append(self.inprogress_constraint.copy(stateful=A_ ) )
								UpperCamelCase						=  None

						if complete:
								# 2. If the next token completes the constraint, move it to completed list, set
								#     inprogress to None. If there are no pending constraints either, then this full list of constraints
								#     is complete.

								self.complete_constraints.append(self.inprogress_constraint )
								UpperCamelCase						=  None

								if len(self.pending_constraints ) == 0:
										# we're done!
										UpperCamelCase						=  True

				else:
						# Not in the middle of fulfilling a constraint. So does this `token_id` helps us step towards any of our list
						# of constraints?

						for cidx, pending_constraint in enumerate(self.pending_constraints ):
								if pending_constraint.does_advance(A_ ):
										UpperCamelCase			,						UpperCamelCase			,						UpperCamelCase						=  pending_constraint.update(A_ )

										if not stepped:
												raise Exception(
												    '`constraint.update(token_id)` is not yielding incremental progress, '
												    'even though `constraint.does_advance(token_id)` is true.' )

										if complete:
												self.complete_constraints.append(A_ )
												UpperCamelCase						=  None

										if not complete and stepped:
												UpperCamelCase						=  pending_constraint

										if complete or stepped:
												# If we made any progress at all, then it's at least not a "pending constraint".

												UpperCamelCase						=  (
												    self.pending_constraints[:cidx] + self.pending_constraints[cidx + 1 :]
												)

												if len(self.pending_constraints ) == 0 and self.inprogress_constraint is None:
														# If there's no longer any pending after this and no inprogress either, then we must be
														# complete.

														UpperCamelCase						=  True

												break  # prevent accidentally stepping through multiple constraints with just one token.

				return complete, stepped



		def __UpperCamelCase				(				self       ,   A_=True )						->      Tuple:
				"""simple docstring"""
				UpperCamelCase						=  ConstraintListState(self.constraints )  # we actually never though self.constraints objects
				# throughout this process. So it's at initialization state.

				if stateful:
						UpperCamelCase						=  [
						    constraint.copy(stateful=A_ ) for constraint in self.complete_constraints
						]
						if self.inprogress_constraint is not None:
								UpperCamelCase						=  self.inprogress_constraint.copy(stateful=A_ )
						UpperCamelCase						=  [constraint.copy() for constraint in self.pending_constraints]

				return new_state
