
'''simple docstring'''







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






class  lowerCAmelCase__   (  lowerCamelCase_     ):






					def __init__(   self		):


											"""simple docstring"""





											self.test()






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     Optional[int]													=       0
											lowercase_	:     Tuple													=       False
											while not completed:
																	if counter == 1:
																							self.reset()
																	lowercase_	:     Union[str, Any]													=       self.advance()
																	if not self.does_advance(__SCREAMING_SNAKE_CASE		):
																							raise Exception(
																							    '''Custom Constraint is not defined correctly. self.does_advance(self.advance()) must be true.'''		)

																	lowercase_						, lowercase_						, lowercase_	:     Dict													=       self.update(__SCREAMING_SNAKE_CASE		)
																	counter += 1

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

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






					@abstractmethod
					def 	_snake_case    (   self		):


											"""simple docstring"""





											raise NotImplementedError(
											    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''		)






					@abstractmethod
					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE		):


											"""simple docstring"""





											raise NotImplementedError(
											    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''		)






					@abstractmethod
					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE		):


											"""simple docstring"""





											raise NotImplementedError(
											    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''		)






					@abstractmethod
					def 	_snake_case    (   self		):


											"""simple docstring"""





											raise NotImplementedError(
											    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''		)






					@abstractmethod
					def 	_snake_case    (   self		):


											"""simple docstring"""





											raise NotImplementedError(
											    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''		)






					@abstractmethod
					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE=False		):


											"""simple docstring"""





											raise NotImplementedError(
											    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''		)







class  lowerCAmelCase__   (  lowerCamelCase_     ):






					def __init__(   self				,							__SCREAMING_SNAKE_CASE		):


											"""simple docstring"""





											super(__SCREAMING_SNAKE_CASE				,							self		).__init__()

											if not isinstance(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		) or len(__SCREAMING_SNAKE_CASE		) == 0:
																	raise ValueError(F'''`token_ids` has to be a non-empty list, but is {token_ids}.'''		)
											if any((not isinstance(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		) 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}.'''		)

											lowercase_	:     List[Any]													=       token_ids

											lowercase_	:     Any													=       len(self.token_ids		)
											lowercase_	:     Optional[Any]													=       -1  # the index of the currently fulfilled step
											lowercase_	:     str													=       False






					def 	_snake_case    (   self		):


											"""simple docstring"""





											if self.completed:
																	return None
											return self.token_ids[self.fulfilled_idx + 1]






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE		):


											"""simple docstring"""





											if not isinstance(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		):
																	raise ValueError(F'''`token_id` has to be an `int`, but is {token_id} of type {type(__SCREAMING_SNAKE_CASE		)}'''		)

											if self.completed:
																	return False

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






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE		):


											"""simple docstring"""





											if not isinstance(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		):
																	raise ValueError(F'''`token_id` has to be an `int`, but is {token_id} of type {type(__SCREAMING_SNAKE_CASE		)}'''		)

											lowercase_	:     str													=       False
											lowercase_	:     Optional[Any]													=       False
											lowercase_	:     Optional[int]													=       False

											if self.does_advance(__SCREAMING_SNAKE_CASE		):
																	self.fulfilled_idx += 1
																	lowercase_	:     List[str]													=       True
																	if self.fulfilled_idx == (self.seqlen - 1):
																							lowercase_	:     List[str]													=       True
																	lowercase_	:     Tuple													=       completed
											else:
																	# failed to make progress.
																	lowercase_	:     Optional[Any]													=       True
																	self.reset()
											return stepped, completed, reset






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     Any													=       False
											lowercase_	:     Union[str, Any]													=       0






					def 	_snake_case    (   self		):


											"""simple docstring"""





											return self.seqlen - (self.fulfilled_idx + 1)






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE=False		):


											"""simple docstring"""





											lowercase_	:     str													=       PhrasalConstraint(self.token_ids		)

											if stateful:
																	lowercase_	:     Optional[int]													=       self.seqlen
																	lowercase_	:     Tuple													=       self.fulfilled_idx
																	lowercase_	:     Dict													=       self.completed

											return new_constraint







class  lowerCAmelCase__   :






					def __init__(   self				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE=True		):


											"""simple docstring"""





											lowercase_	:     Tuple													=       max([len(__SCREAMING_SNAKE_CASE		) for one in nested_token_ids]		)

											lowercase_	:     List[str]													=       {}
											for token_ids in nested_token_ids:
																	lowercase_	:     List[Any]													=       root
																	for tidx, token_id in enumerate(__SCREAMING_SNAKE_CASE		):
																							if token_id not in level:
																													lowercase_	:     Dict													=       {}

																							lowercase_	:     Optional[int]													=       level[token_id]

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

											lowercase_	:     List[str]													=       root






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE		):


											"""simple docstring"""





											lowercase_	:     List[Any]													=       self.trie

											for current_token in current_seq:
																	lowercase_	:     Any													=       start[current_token]

											lowercase_	:     Optional[int]													=       list(start.keys()		)

											return next_tokens






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE		):


											"""simple docstring"""





											lowercase_	:     Tuple													=       self.next_tokens(__SCREAMING_SNAKE_CASE		)

											return len(__SCREAMING_SNAKE_CASE		) == 0






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE		):


											"""simple docstring"""





											lowercase_	:     List[str]													=       list(root.values()		)
											if len(__SCREAMING_SNAKE_CASE		) == 0:
																	return 1
											else:
																	return sum([self.count_leaves(__SCREAMING_SNAKE_CASE		) for nn in next_nodes]		)






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		):


											"""simple docstring"""





											lowercase_	:     List[str]													=       self.count_leaves(__SCREAMING_SNAKE_CASE		)
											return len(__SCREAMING_SNAKE_CASE		) != leaf_count







class  lowerCAmelCase__   (  lowerCamelCase_     ):






					def __init__(   self				,							__SCREAMING_SNAKE_CASE		):


											"""simple docstring"""





											super(__SCREAMING_SNAKE_CASE				,							self		).__init__()

											if not isinstance(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		) or len(__SCREAMING_SNAKE_CASE		) == 0:
																	raise ValueError(F'''`nested_token_ids` has to be a non-empty list, but is {nested_token_ids}.'''		)
											if any(not isinstance(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		) 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(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		) 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}.'''		)

											lowercase_	:     Optional[int]													=       DisjunctiveTrie(__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Dict													=       nested_token_ids

											lowercase_	:     Optional[int]													=       self.trie.max_height
											lowercase_	:     Optional[Any]													=       []
											lowercase_	:     List[Any]													=       False






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     str													=       self.trie.next_tokens(self.current_seq		)

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






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE		):


											"""simple docstring"""





											if not isinstance(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		):
																	raise ValueError(F'''`token_id` is supposed to be type `int`, but is {token_id} of type {type(__SCREAMING_SNAKE_CASE		)}'''		)

											lowercase_	:     Tuple													=       self.trie.next_tokens(self.current_seq		)

											return token_id in next_tokens






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE		):


											"""simple docstring"""





											if not isinstance(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		):
																	raise ValueError(F'''`token_id` is supposed to be type `int`, but is {token_id} of type {type(__SCREAMING_SNAKE_CASE		)}'''		)

											lowercase_	:     Dict													=       False
											lowercase_	:     int													=       False
											lowercase_	:     Dict													=       False

											if self.does_advance(__SCREAMING_SNAKE_CASE		):
																	self.current_seq.append(__SCREAMING_SNAKE_CASE		)
																	lowercase_	:     Optional[int]													=       True
											else:
																	lowercase_	:     List[Any]													=       True
																	self.reset()

											lowercase_	:     List[Any]													=       self.trie.reached_leaf(self.current_seq		)
											lowercase_	:     List[Any]													=       completed

											return stepped, completed, reset






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     int													=       False
											lowercase_	:     Optional[int]													=       []






					def 	_snake_case    (   self		):


											"""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 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE=False		):


											"""simple docstring"""





											lowercase_	:     Tuple													=       DisjunctiveConstraint(self.token_ids		)

											if stateful:
																	lowercase_	:     int													=       self.seqlen
																	lowercase_	:     Tuple													=       self.current_seq
																	lowercase_	:     Optional[Any]													=       self.completed

											return new_constraint







class  lowerCAmelCase__   :






					def __init__(   self				,							__SCREAMING_SNAKE_CASE		):


											"""simple docstring"""





											lowercase_	:     Dict													=       constraints

											# max # of steps required to fulfill a given constraint
											lowercase_	:     str													=       max([c.seqlen for c in constraints]		)
											lowercase_	:     int													=       len(__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Optional[Any]													=       False

											self.init_state()






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     int													=       []
											lowercase_	:     Dict													=       None
											lowercase_	:     Dict													=       [constraint.copy(stateful=__SCREAMING_SNAKE_CASE		) for constraint in self.constraints]






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     Dict													=       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 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     List[str]													=       []
											if self.inprogress_constraint is None:
																	for constraint in self.pending_constraints:  # "pending" == "unfulfilled yet"
																							lowercase_	:     Union[str, Any]													=       constraint.advance()
																							if isinstance(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		):
																													token_list.append(__SCREAMING_SNAKE_CASE		)
																							elif isinstance(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		):
																													token_list.extend(__SCREAMING_SNAKE_CASE		)
											else:
																	lowercase_	:     str													=       self.inprogress_constraint.advance()
																	if isinstance(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		):
																							token_list.append(__SCREAMING_SNAKE_CASE		)
																	elif isinstance(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		):
																							token_list.extend(__SCREAMING_SNAKE_CASE		)

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






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE		):


											"""simple docstring"""





											self.init_state()

											if token_ids is not None:
																	for token in token_ids:
																							# completes or steps **one** constraint
																							lowercase_						, lowercase_	:     List[Any]													=       self.add(__SCREAMING_SNAKE_CASE		)

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






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE		):


											"""simple docstring"""





											if not isinstance(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		):
																	raise ValueError(F'''`token_id` should be an `int`, but is `{token_id}`.'''		)

											lowercase_						, lowercase_	:     List[str]													=       False, False

											if self.completed:
																	lowercase_	:     List[str]													=       True
																	lowercase_	:     Tuple													=       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

																	lowercase_						, lowercase_						, lowercase_	:     Tuple													=       self.inprogress_constraint.update(__SCREAMING_SNAKE_CASE		)
																	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=__SCREAMING_SNAKE_CASE		)		)
																							lowercase_	:     Tuple													=       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		)
																							lowercase_	:     List[Any]													=       None

																							if len(self.pending_constraints		) == 0:
																													# we're done!
																													lowercase_	:     int													=       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(__SCREAMING_SNAKE_CASE		):
																													lowercase_						, lowercase_						, lowercase_	:     str													=       pending_constraint.update(__SCREAMING_SNAKE_CASE		)

																													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(__SCREAMING_SNAKE_CASE		)
																																			lowercase_	:     str													=       None

																													if not complete and stepped:
																																			lowercase_	:     Any													=       pending_constraint

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

																																			lowercase_	:     List[Any]													=       (
																																			    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.

																																									lowercase_	:     Union[str, Any]													=       True

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

											return complete, stepped






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE=True		):


											"""simple docstring"""





											lowercase_	:     str													=       ConstraintListState(self.constraints		)  # we actually never though self.constraints objects
											# throughout this process. So it's at initialization state.

											if stateful:
																	lowercase_	:     Dict													=       [
																	    constraint.copy(stateful=__SCREAMING_SNAKE_CASE		) for constraint in self.complete_constraints
																	]
																	if self.inprogress_constraint is not None:
																							lowercase_	:     int													=       self.inprogress_constraint.copy(stateful=__SCREAMING_SNAKE_CASE		)
																	lowercase_	:     Union[str, Any]													=       [constraint.copy() for constraint in self.pending_constraints]

											return new_state





'''simple docstring'''







from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


_lowercase :  int    							=       logging.get_logger(__name__)

_lowercase :  List[Any]    							=       {
    "shi-labs/nat-mini-in1k-224": "https://huggingface.co/shi-labs/nat-mini-in1k-224/resolve/main/config.json",
    # See all Nat models at https://huggingface.co/models?filter=nat
}






class  lowerCAmelCase__   (  lowerCamelCase_    ,						lowerCamelCase_     ):
					lowerCAmelCase_        =     '''nat'''

					lowerCAmelCase_        =     {
					    '''num_attention_heads''': '''num_heads''',
					    '''num_hidden_layers''': '''num_layers''',
					}






					def __init__(   self				,							__SCREAMING_SNAKE_CASE=4				,							__SCREAMING_SNAKE_CASE=3				,							__SCREAMING_SNAKE_CASE=64				,							__SCREAMING_SNAKE_CASE=[3, 4, 6, 5]				,							__SCREAMING_SNAKE_CASE=[2, 4, 8, 16]				,							__SCREAMING_SNAKE_CASE=7				,							__SCREAMING_SNAKE_CASE=3.0				,							__SCREAMING_SNAKE_CASE=True				,							__SCREAMING_SNAKE_CASE=0.0				,							__SCREAMING_SNAKE_CASE=0.0				,							__SCREAMING_SNAKE_CASE=0.1				,							__SCREAMING_SNAKE_CASE="gelu"				,							__SCREAMING_SNAKE_CASE=0.02				,							__SCREAMING_SNAKE_CASE=1E-5				,							__SCREAMING_SNAKE_CASE=0.0				,							__SCREAMING_SNAKE_CASE=None				,							__SCREAMING_SNAKE_CASE=None				,							**__SCREAMING_SNAKE_CASE				,							):


											"""simple docstring"""





											super().__init__(**__SCREAMING_SNAKE_CASE		)

											lowercase_	:     Optional[int]													=       patch_size
											lowercase_	:     List[Any]													=       num_channels
											lowercase_	:     str													=       embed_dim
											lowercase_	:     List[str]													=       depths
											lowercase_	:     str													=       len(__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Tuple													=       num_heads
											lowercase_	:     int													=       kernel_size
											lowercase_	:     Union[str, Any]													=       mlp_ratio
											lowercase_	:     Optional[int]													=       qkv_bias
											lowercase_	:     List[Any]													=       hidden_dropout_prob
											lowercase_	:     Optional[int]													=       attention_probs_dropout_prob
											lowercase_	:     List[Any]													=       drop_path_rate
											lowercase_	:     List[Any]													=       hidden_act
											lowercase_	:     int													=       layer_norm_eps
											lowercase_	:     int													=       initializer_range
											# we set the hidden_size attribute in order to make Nat work with VisionEncoderDecoderModel
											# this indicates the channel dimension after the last stage of the model
											lowercase_	:     Dict													=       int(embed_dim * 2 ** (len(__SCREAMING_SNAKE_CASE		) - 1)		)
											lowercase_	:     Tuple													=       layer_scale_init_value
											lowercase_	:     Union[str, Any]													=       ['''stem'''] + [F'''stage{idx}''' for idx in range(1				,							len(__SCREAMING_SNAKE_CASE		) + 1		)]
											lowercase_						, lowercase_	:     int													=       get_aligned_output_features_output_indices(
											    out_features=__SCREAMING_SNAKE_CASE				,							out_indices=__SCREAMING_SNAKE_CASE				,							stage_names=self.stage_names		)



