
import argparse
import os

import gluonnlp as nlp
import mxnet as mx
import numpy as np
import torch
from gluonnlp.base import get_home_dir
from gluonnlp.model.bert import BERTEncoder
from gluonnlp.model.utils import _load_vocab
from gluonnlp.vocab import Vocab
from packaging import version
from torch import nn

from transformers import BertConfig, BertForMaskedLM, BertModel, RobertaTokenizer
from transformers.models.bert.modeling_bert import (
    BertIntermediate,
    BertLayer,
    BertOutput,
    BertSelfAttention,
    BertSelfOutput,
)
from transformers.utils import logging


if version.parse(nlp.__version__) != version.parse("0.8.3"):
			raise Exception("requires gluonnlp == 0.8.3")

if version.parse(mx.__version__) != version.parse("1.5.0"):
			raise Exception("requires mxnet == 1.5.0")

logging.set_verbosity_info()
_UpperCAmelCase      :				Tuple          = logging.get_logger(__name__)

_UpperCAmelCase      :				List[str]          = """The Nymphenburg Palace is a beautiful palace in Munich!"""


def    A (      lowercase					, lowercase				)     ->						int:
		'''simple docstring'''

		UpperCamelCase						=  {
		    'attention_cell': 'multi_head',
		    'num_layers': 4,
		    'units': 1_024,
		    'hidden_size': 768,
		    'max_length': 512,
		    'num_heads': 8,
		    'scaled': True,
		    'dropout': 0.1,
		    'use_residual': True,
		    'embed_size': 1_024,
		    'embed_dropout': 0.1,
		    'word_embed': None,
		    'layer_norm_eps': 1e-5,
		    'token_type_vocab_size': 2,
		}

		UpperCamelCase						=  bort_4_8_768_1024_hparams

		# Let's construct the original Bort model here
		# Taken from official BERT implementation, see:
		# https://github.com/alexa/bort/blob/master/bort/bort.py
		UpperCamelCase						=  BERTEncoder(
		    attention_cell=predefined_args['attention_cell']					, num_layers=predefined_args['num_layers']					, units=predefined_args['units']					, hidden_size=predefined_args['hidden_size']					, max_length=predefined_args['max_length']					, num_heads=predefined_args['num_heads']					, scaled=predefined_args['scaled']					, dropout=predefined_args['dropout']					, output_attention=lowercase					, output_all_encodings=lowercase					, use_residual=predefined_args['use_residual']					, activation=predefined_args.get('activation'					, 'gelu'				)					, layer_norm_eps=predefined_args.get('layer_norm_eps'					, lowercase				)					, )

		# Vocab information needs to be fetched first
		# It's the same as RoBERTa, so RobertaTokenizer can be used later
		UpperCamelCase						=  'openwebtext_ccnews_stories_books_cased'

		# Specify download folder to Gluonnlp's vocab
		UpperCamelCase						=  os.path.join(get_home_dir()					, 'models'				)
		UpperCamelCase						=  _load_vocab(lowercase					, lowercase					, lowercase					, cls=lowercase				)

		UpperCamelCase						=  nlp.model.BERTModel(
		    lowercase					, len(lowercase				)					, units=predefined_args['units']					, embed_size=predefined_args['embed_size']					, embed_dropout=predefined_args['embed_dropout']					, word_embed=predefined_args['word_embed']					, use_pooler=lowercase					, use_token_type_embed=lowercase					, token_type_vocab_size=predefined_args['token_type_vocab_size']					, use_classifier=lowercase					, use_decoder=lowercase					, )

		original_bort.load_parameters(lowercase					, cast_dtype=lowercase					, ignore_extra=lowercase				)
		UpperCamelCase						=  original_bort._collect_params_with_prefix()

		# Build our config 🤗
		UpperCamelCase						=  {
		    'architectures': ['BertForMaskedLM'],
		    'attention_probs_dropout_prob': predefined_args['dropout'],
		    'hidden_act': 'gelu',
		    'hidden_dropout_prob': predefined_args['dropout'],
		    'hidden_size': predefined_args['embed_size'],
		    'initializer_range': 0.0_2,
		    'intermediate_size': predefined_args['hidden_size'],
		    'layer_norm_eps': predefined_args['layer_norm_eps'],
		    'max_position_embeddings': predefined_args['max_length'],
		    'model_type': 'bort',
		    'num_attention_heads': predefined_args['num_heads'],
		    'num_hidden_layers': predefined_args['num_layers'],
		    'pad_token_id': 1,  # 2 = BERT, 1 = RoBERTa
		    'type_vocab_size': 1,  # 2 = BERT, 1 = RoBERTa
		    'vocab_size': len(lowercase				),
		}

		UpperCamelCase						=  BertConfig.from_dict(lowercase				)
		UpperCamelCase						=  BertForMaskedLM(lowercase				)
		hf_bort_model.eval()

		# Parameter mapping table (Gluonnlp to Transformers)
		# * denotes layer index
		#
		# | Gluon Parameter                                                | Transformers Parameter
		# | -------------------------------------------------------------- | ----------------------
		# | `encoder.layer_norm.beta`                                      | `bert.embeddings.LayerNorm.bias`
		# | `encoder.layer_norm.gamma`                                     | `bert.embeddings.LayerNorm.weight`
		# | `encoder.position_weight`                                      | `bert.embeddings.position_embeddings.weight`
		# | `word_embed.0.weight`                                          | `bert.embeddings.word_embeddings.weight`
		# | `encoder.transformer_cells.*.attention_cell.proj_key.bias`     | `bert.encoder.layer.*.attention.self.key.bias`
		# | `encoder.transformer_cells.*.attention_cell.proj_key.weight`   | `bert.encoder.layer.*.attention.self.key.weight`
		# | `encoder.transformer_cells.*.attention_cell.proj_query.bias`   | `bert.encoder.layer.*.attention.self.query.bias`
		# | `encoder.transformer_cells.*.attention_cell.proj_query.weight` | `bert.encoder.layer.*.attention.self.query.weight`
		# | `encoder.transformer_cells.*.attention_cell.proj_value.bias`   | `bert.encoder.layer.*.attention.self.value.bias`
		# | `encoder.transformer_cells.*.attention_cell.proj_value.weight` | `bert.encoder.layer.*.attention.self.value.weight`
		# | `encoder.transformer_cells.*.ffn.ffn_2.bias`                   | `bert.encoder.layer.*.attention.output.dense.bias`
		# | `encoder.transformer_cells.*.ffn.ffn_2.weight`                 | `bert.encoder.layer.*.attention.output.dense.weight`
		# | `encoder.transformer_cells.*.layer_norm.beta`                  | `bert.encoder.layer.*.attention.output.LayerNorm.bias`
		# | `encoder.transformer_cells.*.layer_norm.gamma`                 | `bert.encoder.layer.*.attention.output.LayerNorm.weight`
		# | `encoder.transformer_cells.*.ffn.ffn_1.bias`                   | `bert.encoder.layer.*.intermediate.dense.bias`
		# | `encoder.transformer_cells.*.ffn.ffn_1.weight`                 | `bert.encoder.layer.*.intermediate.dense.weight`
		# | `encoder.transformer_cells.*.ffn.layer_norm.beta`              | `bert.encoder.layer.*.output.LayerNorm.bias`
		# | `encoder.transformer_cells.*.ffn.layer_norm.gamma`             | `bert.encoder.layer.*.output.LayerNorm.weight`
		# | `encoder.transformer_cells.*.proj.bias`                        | `bert.encoder.layer.*.output.dense.bias`
		# | `encoder.transformer_cells.*.proj.weight`                      | `bert.encoder.layer.*.output.dense.weight`

		# Helper function to convert MXNET Arrays to PyTorch
		def to_torch(lowercase				) -> nn.Parameter:
				return nn.Parameter(torch.FloatTensor(mx_array.data().asnumpy()				)				)

		# Check param shapes and map new HF param back
		def check_and_map_params(lowercase					, lowercase				):
				UpperCamelCase						=  hf_param.shape

				UpperCamelCase						=  to_torch(params[gluon_param]				)
				UpperCamelCase						=  gluon_param.shape

				assert (
				    shape_hf == shape_gluon
				), f'''The gluon parameter {gluon_param} has shape {shape_gluon}, but expects shape {shape_hf} for Transformers'''

				return gluon_param

		UpperCamelCase						=  check_and_map_params(
		    hf_bort_model.bert.embeddings.word_embeddings.weight					, 'word_embed.0.weight'				)
		UpperCamelCase						=  check_and_map_params(
		    hf_bort_model.bert.embeddings.position_embeddings.weight					, 'encoder.position_weight'				)
		UpperCamelCase						=  check_and_map_params(
		    hf_bort_model.bert.embeddings.LayerNorm.bias					, 'encoder.layer_norm.beta'				)
		UpperCamelCase						=  check_and_map_params(
		    hf_bort_model.bert.embeddings.LayerNorm.weight					, 'encoder.layer_norm.gamma'				)

		# Inspired by RoBERTa conversion script, we just zero them out (Bort does not use them)
		UpperCamelCase						=  torch.zeros_like(
		    hf_bort_model.bert.embeddings.token_type_embeddings.weight.data				)

		for i in range(hf_bort_config.num_hidden_layers				):
				UpperCamelCase						=  hf_bort_model.bert.encoder.layer[i]

				# self attention
				UpperCamelCase						=  layer.attention.self

				UpperCamelCase						=  check_and_map_params(
				    self_attn.key.bias.data					, f'''encoder.transformer_cells.{i}.attention_cell.proj_key.bias'''				)

				UpperCamelCase						=  check_and_map_params(
				    self_attn.key.weight.data					, f'''encoder.transformer_cells.{i}.attention_cell.proj_key.weight'''				)
				UpperCamelCase						=  check_and_map_params(
				    self_attn.query.bias.data					, f'''encoder.transformer_cells.{i}.attention_cell.proj_query.bias'''				)
				UpperCamelCase						=  check_and_map_params(
				    self_attn.query.weight.data					, f'''encoder.transformer_cells.{i}.attention_cell.proj_query.weight'''				)
				UpperCamelCase						=  check_and_map_params(
				    self_attn.value.bias.data					, f'''encoder.transformer_cells.{i}.attention_cell.proj_value.bias'''				)
				UpperCamelCase						=  check_and_map_params(
				    self_attn.value.weight.data					, f'''encoder.transformer_cells.{i}.attention_cell.proj_value.weight'''				)

				# self attention output
				UpperCamelCase						=  layer.attention.output

				UpperCamelCase						=  check_and_map_params(
				    self_output.dense.bias					, f'''encoder.transformer_cells.{i}.proj.bias'''				)
				UpperCamelCase						=  check_and_map_params(
				    self_output.dense.weight					, f'''encoder.transformer_cells.{i}.proj.weight'''				)
				UpperCamelCase						=  check_and_map_params(
				    self_output.LayerNorm.bias					, f'''encoder.transformer_cells.{i}.layer_norm.beta'''				)
				UpperCamelCase						=  check_and_map_params(
				    self_output.LayerNorm.weight					, f'''encoder.transformer_cells.{i}.layer_norm.gamma'''				)

				# intermediate
				UpperCamelCase						=  layer.intermediate

				UpperCamelCase						=  check_and_map_params(
				    intermediate.dense.bias					, f'''encoder.transformer_cells.{i}.ffn.ffn_1.bias'''				)
				UpperCamelCase						=  check_and_map_params(
				    intermediate.dense.weight					, f'''encoder.transformer_cells.{i}.ffn.ffn_1.weight'''				)

				# output
				UpperCamelCase						=  layer.output

				UpperCamelCase						=  check_and_map_params(
				    bert_output.dense.bias					, f'''encoder.transformer_cells.{i}.ffn.ffn_2.bias'''				)
				UpperCamelCase						=  check_and_map_params(
				    bert_output.dense.weight					, f'''encoder.transformer_cells.{i}.ffn.ffn_2.weight'''				)
				UpperCamelCase						=  check_and_map_params(
				    bert_output.LayerNorm.bias					, f'''encoder.transformer_cells.{i}.ffn.layer_norm.beta'''				)
				UpperCamelCase						=  check_and_map_params(
				    bert_output.LayerNorm.weight					, f'''encoder.transformer_cells.{i}.ffn.layer_norm.gamma'''				)

		# Save space and energy 🎄
		hf_bort_model.half()

		# Compare output of both models
		UpperCamelCase						=  RobertaTokenizer.from_pretrained('roberta-base'				)

		UpperCamelCase						=  tokenizer.encode_plus(lowercase				)['input_ids']

		# Get gluon output
		UpperCamelCase						=  mx.nd.array([input_ids]				)
		UpperCamelCase						=  original_bort(inputs=lowercase					, token_types=[]				)

		# Get Transformer output (save and reload model again)
		hf_bort_model.save_pretrained(lowercase				)
		UpperCamelCase						=  BertModel.from_pretrained(lowercase				)
		hf_bort_model.eval()

		UpperCamelCase						=  tokenizer.encode_plus(lowercase					, return_tensors='pt'				)
		UpperCamelCase						=  hf_bort_model(**lowercase				)[0]

		UpperCamelCase						=  output_gluon[0].asnumpy()
		UpperCamelCase						=  output_hf[0].detach().numpy()

		UpperCamelCase						=  np.max(np.abs(hf_layer - gluon_layer				)				).item()
		UpperCamelCase						=  np.allclose(lowercase					, lowercase					, atol=1e-3				)

		if success:
				print('✔️ Both model do output the same tensors'				)
		else:
				print('❌ Both model do **NOT** output the same tensors'				)
				print('Absolute difference is:'					, lowercase				)


if __name__ == "__main__":
			_UpperCAmelCase      :				str          = argparse.ArgumentParser()
			# Required parameters
			parser.add_argument(
			    "--bort_checkpoint_path", default=None, type=str, required=True, help="Path the official Bort params file."
			)
			parser.add_argument(
			    "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
			)
			_UpperCAmelCase      :				List[Any]          = parser.parse_args()
			convert_bort_checkpoint_to_pytorch(args.bort_checkpoint_path, args.pytorch_dump_folder_path)


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
