



'''simple docstring'''



import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import cached_download, hf_hub_download, hf_hub_url
from PIL import Image

from transformers import DetaConfig, DetaForObjectDetection, DetaImageProcessor, SwinConfig
from transformers.utils import logging


logging.set_verbosity_info()
lowerCAmelCase	:		int				=   logging.get_logger(__name__)



def   __lowerCAmelCase					(							lowerCamelCase   :						int					):






	'''simple docstring'''




	__lowerCAmelCase												=  SwinConfig(
	    embed_dim=1_92				,		depths=(2, 2, 18, 2)				,		num_heads=(6, 12, 24, 48)				,		window_size=12				,		out_features=["stage2", "stage3", "stage4"]				,		)

	__lowerCAmelCase												=  DetaConfig(
	    backbone_config=lowerCamelCase				,		num_queries=9_00				,		encoder_ffn_dim=20_48				,		decoder_ffn_dim=20_48				,		num_feature_levels=5				,		assign_first_stage=lowerCamelCase				,		with_box_refine=lowerCamelCase				,		two_stage=lowerCamelCase				,		)

	# set labels
	__lowerCAmelCase												=  "huggingface/label-files"
	if "o365" in model_name:
		__lowerCAmelCase												=  3_66
		__lowerCAmelCase												=  "object365-id2label.json"
	else:
		__lowerCAmelCase												=  91
		__lowerCAmelCase												=  "coco-detection-id2label.json"

	__lowerCAmelCase												=  num_labels
	__lowerCAmelCase												=  json.load(open(cached_download(hf_hub_url(lowerCamelCase				,		lowerCamelCase				,		repo_type="dataset"					)					)				,		"r"					)					)
	__lowerCAmelCase												=  {int(lowerCamelCase					): v for k, v in idalabel.items()}
	__lowerCAmelCase												=  idalabel
	__lowerCAmelCase												=  {v: k for k, v in idalabel.items()}

	return config



def   __lowerCAmelCase					(							lowerCamelCase   :						Dict					):






	'''simple docstring'''




	__lowerCAmelCase												=  []

	# stem
	# fmt: off
	rename_keys.append(("backbone.0.body.patch_embed.proj.weight", "model.backbone.model.embeddings.patch_embeddings.projection.weight")					)
	rename_keys.append(("backbone.0.body.patch_embed.proj.bias", "model.backbone.model.embeddings.patch_embeddings.projection.bias")					)
	rename_keys.append(("backbone.0.body.patch_embed.norm.weight", "model.backbone.model.embeddings.norm.weight")					)
	rename_keys.append(("backbone.0.body.patch_embed.norm.bias", "model.backbone.model.embeddings.norm.bias")					)
	# stages
	for i in range(len(config.backbone_config.depths					)					):
		for j in range(config.backbone_config.depths[i]					):
			rename_keys.append((f'''backbone.0.body.layers.{i}.blocks.{j}.norm1.weight''', f'''model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_before.weight''')					)
			rename_keys.append((f'''backbone.0.body.layers.{i}.blocks.{j}.norm1.bias''', f'''model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_before.bias''')					)
			rename_keys.append((f'''backbone.0.body.layers.{i}.blocks.{j}.attn.relative_position_bias_table''', f'''model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_bias_table''')					)
			rename_keys.append((f'''backbone.0.body.layers.{i}.blocks.{j}.attn.relative_position_index''', f'''model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_index''')					)
			rename_keys.append((f'''backbone.0.body.layers.{i}.blocks.{j}.attn.proj.weight''', f'''model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.output.dense.weight''')					)
			rename_keys.append((f'''backbone.0.body.layers.{i}.blocks.{j}.attn.proj.bias''', f'''model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.output.dense.bias''')					)
			rename_keys.append((f'''backbone.0.body.layers.{i}.blocks.{j}.norm2.weight''', f'''model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_after.weight''')					)
			rename_keys.append((f'''backbone.0.body.layers.{i}.blocks.{j}.norm2.bias''', f'''model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_after.bias''')					)
			rename_keys.append((f'''backbone.0.body.layers.{i}.blocks.{j}.mlp.fc1.weight''', f'''model.backbone.model.encoder.layers.{i}.blocks.{j}.intermediate.dense.weight''')					)
			rename_keys.append((f'''backbone.0.body.layers.{i}.blocks.{j}.mlp.fc1.bias''', f'''model.backbone.model.encoder.layers.{i}.blocks.{j}.intermediate.dense.bias''')					)
			rename_keys.append((f'''backbone.0.body.layers.{i}.blocks.{j}.mlp.fc2.weight''', f'''model.backbone.model.encoder.layers.{i}.blocks.{j}.output.dense.weight''')					)
			rename_keys.append((f'''backbone.0.body.layers.{i}.blocks.{j}.mlp.fc2.bias''', f'''model.backbone.model.encoder.layers.{i}.blocks.{j}.output.dense.bias''')					)

		if i < 3:
			rename_keys.append((f'''backbone.0.body.layers.{i}.downsample.reduction.weight''', f'''model.backbone.model.encoder.layers.{i}.downsample.reduction.weight''')					)
			rename_keys.append((f'''backbone.0.body.layers.{i}.downsample.norm.weight''', f'''model.backbone.model.encoder.layers.{i}.downsample.norm.weight''')					)
			rename_keys.append((f'''backbone.0.body.layers.{i}.downsample.norm.bias''', f'''model.backbone.model.encoder.layers.{i}.downsample.norm.bias''')					)

	rename_keys.append(("backbone.0.body.norm1.weight", "model.backbone.model.hidden_states_norms.stage2.weight")					)
	rename_keys.append(("backbone.0.body.norm1.bias", "model.backbone.model.hidden_states_norms.stage2.bias")					)
	rename_keys.append(("backbone.0.body.norm2.weight", "model.backbone.model.hidden_states_norms.stage3.weight")					)
	rename_keys.append(("backbone.0.body.norm2.bias", "model.backbone.model.hidden_states_norms.stage3.bias")					)
	rename_keys.append(("backbone.0.body.norm3.weight", "model.backbone.model.hidden_states_norms.stage4.weight")					)
	rename_keys.append(("backbone.0.body.norm3.bias", "model.backbone.model.hidden_states_norms.stage4.bias")					)

	# transformer encoder
	for i in range(config.encoder_layers					):
		rename_keys.append((f'''transformer.encoder.layers.{i}.self_attn.sampling_offsets.weight''', f'''model.encoder.layers.{i}.self_attn.sampling_offsets.weight''')					)
		rename_keys.append((f'''transformer.encoder.layers.{i}.self_attn.sampling_offsets.bias''', f'''model.encoder.layers.{i}.self_attn.sampling_offsets.bias''')					)
		rename_keys.append((f'''transformer.encoder.layers.{i}.self_attn.attention_weights.weight''', f'''model.encoder.layers.{i}.self_attn.attention_weights.weight''')					)
		rename_keys.append((f'''transformer.encoder.layers.{i}.self_attn.attention_weights.bias''', f'''model.encoder.layers.{i}.self_attn.attention_weights.bias''')					)
		rename_keys.append((f'''transformer.encoder.layers.{i}.self_attn.value_proj.weight''', f'''model.encoder.layers.{i}.self_attn.value_proj.weight''')					)
		rename_keys.append((f'''transformer.encoder.layers.{i}.self_attn.value_proj.bias''', f'''model.encoder.layers.{i}.self_attn.value_proj.bias''')					)
		rename_keys.append((f'''transformer.encoder.layers.{i}.self_attn.output_proj.weight''', f'''model.encoder.layers.{i}.self_attn.output_proj.weight''')					)
		rename_keys.append((f'''transformer.encoder.layers.{i}.self_attn.output_proj.bias''', f'''model.encoder.layers.{i}.self_attn.output_proj.bias''')					)
		rename_keys.append((f'''transformer.encoder.layers.{i}.norm1.weight''', f'''model.encoder.layers.{i}.self_attn_layer_norm.weight''')					)
		rename_keys.append((f'''transformer.encoder.layers.{i}.norm1.bias''', f'''model.encoder.layers.{i}.self_attn_layer_norm.bias''')					)
		rename_keys.append((f'''transformer.encoder.layers.{i}.linear1.weight''', f'''model.encoder.layers.{i}.fc1.weight''')					)
		rename_keys.append((f'''transformer.encoder.layers.{i}.linear1.bias''', f'''model.encoder.layers.{i}.fc1.bias''')					)
		rename_keys.append((f'''transformer.encoder.layers.{i}.linear2.weight''', f'''model.encoder.layers.{i}.fc2.weight''')					)
		rename_keys.append((f'''transformer.encoder.layers.{i}.linear2.bias''', f'''model.encoder.layers.{i}.fc2.bias''')					)
		rename_keys.append((f'''transformer.encoder.layers.{i}.norm2.weight''', f'''model.encoder.layers.{i}.final_layer_norm.weight''')					)
		rename_keys.append((f'''transformer.encoder.layers.{i}.norm2.bias''', f'''model.encoder.layers.{i}.final_layer_norm.bias''')					)

	# transformer decoder
	for i in range(config.decoder_layers					):
		rename_keys.append((f'''transformer.decoder.layers.{i}.cross_attn.sampling_offsets.weight''', f'''model.decoder.layers.{i}.encoder_attn.sampling_offsets.weight''')					)
		rename_keys.append((f'''transformer.decoder.layers.{i}.cross_attn.sampling_offsets.bias''', f'''model.decoder.layers.{i}.encoder_attn.sampling_offsets.bias''')					)
		rename_keys.append((f'''transformer.decoder.layers.{i}.cross_attn.attention_weights.weight''', f'''model.decoder.layers.{i}.encoder_attn.attention_weights.weight''')					)
		rename_keys.append((f'''transformer.decoder.layers.{i}.cross_attn.attention_weights.bias''', f'''model.decoder.layers.{i}.encoder_attn.attention_weights.bias''')					)
		rename_keys.append((f'''transformer.decoder.layers.{i}.cross_attn.value_proj.weight''', f'''model.decoder.layers.{i}.encoder_attn.value_proj.weight''')					)
		rename_keys.append((f'''transformer.decoder.layers.{i}.cross_attn.value_proj.bias''', f'''model.decoder.layers.{i}.encoder_attn.value_proj.bias''')					)
		rename_keys.append((f'''transformer.decoder.layers.{i}.cross_attn.output_proj.weight''', f'''model.decoder.layers.{i}.encoder_attn.output_proj.weight''')					)
		rename_keys.append((f'''transformer.decoder.layers.{i}.cross_attn.output_proj.bias''', f'''model.decoder.layers.{i}.encoder_attn.output_proj.bias''')					)
		rename_keys.append((f'''transformer.decoder.layers.{i}.norm1.weight''', f'''model.decoder.layers.{i}.encoder_attn_layer_norm.weight''')					)
		rename_keys.append((f'''transformer.decoder.layers.{i}.norm1.bias''', f'''model.decoder.layers.{i}.encoder_attn_layer_norm.bias''')					)
		rename_keys.append((f'''transformer.decoder.layers.{i}.self_attn.out_proj.weight''', f'''model.decoder.layers.{i}.self_attn.out_proj.weight''')					)
		rename_keys.append((f'''transformer.decoder.layers.{i}.self_attn.out_proj.bias''', f'''model.decoder.layers.{i}.self_attn.out_proj.bias''')					)
		rename_keys.append((f'''transformer.decoder.layers.{i}.norm2.weight''', f'''model.decoder.layers.{i}.self_attn_layer_norm.weight''')					)
		rename_keys.append((f'''transformer.decoder.layers.{i}.norm2.bias''', f'''model.decoder.layers.{i}.self_attn_layer_norm.bias''')					)
		rename_keys.append((f'''transformer.decoder.layers.{i}.linear1.weight''', f'''model.decoder.layers.{i}.fc1.weight''')					)
		rename_keys.append((f'''transformer.decoder.layers.{i}.linear1.bias''', f'''model.decoder.layers.{i}.fc1.bias''')					)
		rename_keys.append((f'''transformer.decoder.layers.{i}.linear2.weight''', f'''model.decoder.layers.{i}.fc2.weight''')					)
		rename_keys.append((f'''transformer.decoder.layers.{i}.linear2.bias''', f'''model.decoder.layers.{i}.fc2.bias''')					)
		rename_keys.append((f'''transformer.decoder.layers.{i}.norm3.weight''', f'''model.decoder.layers.{i}.final_layer_norm.weight''')					)
		rename_keys.append((f'''transformer.decoder.layers.{i}.norm3.bias''', f'''model.decoder.layers.{i}.final_layer_norm.bias''')					)

	# fmt: on

	return rename_keys



def   __lowerCAmelCase					(							lowerCamelCase   :						List[str]				,		lowerCamelCase   :						Optional[int]				,		lowerCamelCase   :						Dict					):






	'''simple docstring'''




	__lowerCAmelCase												=  dct.pop(lowerCamelCase					)
	__lowerCAmelCase												=  val



def   __lowerCAmelCase					(							lowerCamelCase   :						List[Any]				,		lowerCamelCase   :						Dict					):






	'''simple docstring'''




	__lowerCAmelCase												=  [int(backbone_config.embed_dim * 2**i					) for i in range(len(backbone_config.depths					)					)]
	for i in range(len(backbone_config.depths					)					):
		__lowerCAmelCase												=  num_features[i]
		for j in range(backbone_config.depths[i]					):
			# fmt: off
			# read in weights + bias of input projection layer (in original implementation, this is a single matrix + bias)
			__lowerCAmelCase												=  state_dict.pop(f'''backbone.0.body.layers.{i}.blocks.{j}.attn.qkv.weight'''					)
			__lowerCAmelCase												=  state_dict.pop(f'''backbone.0.body.layers.{i}.blocks.{j}.attn.qkv.bias'''					)
			# next, add query, keys and values (in that order) to the state dict
			__lowerCAmelCase												=  in_proj_weight[:dim, :]
			__lowerCAmelCase												=  in_proj_bias[: dim]
			__lowerCAmelCase												=  in_proj_weight[
			    dim : dim * 2, :
			]
			__lowerCAmelCase												=  in_proj_bias[
			    dim : dim * 2
			]
			__lowerCAmelCase												=  in_proj_weight[
			    -dim :, :
			]
			__lowerCAmelCase												=  in_proj_bias[-dim :]
			# fmt: on



def   __lowerCAmelCase					(							lowerCamelCase   :						Dict				,		lowerCamelCase   :						int					):






	'''simple docstring'''




	__lowerCAmelCase												=  config.d_model
	for i in range(config.decoder_layers					):
		# read in weights + bias of input projection layer of self-attention
		__lowerCAmelCase												=  state_dict.pop(f'''transformer.decoder.layers.{i}.self_attn.in_proj_weight'''					)
		__lowerCAmelCase												=  state_dict.pop(f'''transformer.decoder.layers.{i}.self_attn.in_proj_bias'''					)
		# next, add query, keys and values (in that order) to the state dict
		__lowerCAmelCase												=  in_proj_weight[:hidden_size, :]
		__lowerCAmelCase												=  in_proj_bias[:hidden_size]
		__lowerCAmelCase												=  in_proj_weight[
		    hidden_size : hidden_size * 2, :
		]
		__lowerCAmelCase												=  in_proj_bias[hidden_size : hidden_size * 2]
		__lowerCAmelCase												=  in_proj_weight[-hidden_size:, :]
		__lowerCAmelCase												=  in_proj_bias[-hidden_size:]



def   __lowerCAmelCase					(							):






	'''simple docstring'''




	__lowerCAmelCase												=  "http://images.cocodataset.org/val2017/000000039769.jpg"
	__lowerCAmelCase												=  Image.open(requests.get(lowerCamelCase				,		stream=lowerCamelCase					).raw					)

	return im



@torch.no_grad()
def   __lowerCAmelCase					(							lowerCamelCase   :						List[Any]				,		lowerCamelCase   :						Union[str, Any]				,		lowerCamelCase   :						Any					):






	'''simple docstring'''




	__lowerCAmelCase												=  get_deta_config(lowerCamelCase					)

	# load original state dict
	if model_name == "deta-swin-large":
		__lowerCAmelCase												=  hf_hub_download(repo_id="nielsr/deta-checkpoints"				,		filename="adet_swin_ft.pth"					)
	elif model_name == "deta-swin-large-o365":
		__lowerCAmelCase												=  hf_hub_download(repo_id="jozhang97/deta-swin-l-o365"				,		filename="deta_swin_pt_o365.pth"					)
	else:
		raise ValueError(f'''Model name {model_name} not supported'''					)

	__lowerCAmelCase												=  torch.load(lowerCamelCase				,		map_location="cpu"					)["model"]

	# original state dict
	for name, param in state_dict.items():
		print(lowerCamelCase				,		param.shape					)

	# rename keys
	__lowerCAmelCase												=  create_rename_keys(lowerCamelCase					)
	for src, dest in rename_keys:
		rename_key(lowerCamelCase				,		lowerCamelCase				,		lowerCamelCase					)
	read_in_swin_q_k_v(lowerCamelCase				,		config.backbone_config					)
	read_in_decoder_q_k_v(lowerCamelCase				,		lowerCamelCase					)

	# fix some prefixes
	for key in state_dict.copy().keys():
		if "transformer.decoder.class_embed" in key or "transformer.decoder.bbox_embed" in key:
			__lowerCAmelCase												=  state_dict.pop(lowerCamelCase					)
			__lowerCAmelCase												=  val
		if "input_proj" in key:
			__lowerCAmelCase												=  state_dict.pop(lowerCamelCase					)
			__lowerCAmelCase												=  val
		if "level_embed" in key or "pos_trans" in key or "pix_trans" in key or "enc_output" in key:
			__lowerCAmelCase												=  state_dict.pop(lowerCamelCase					)
			__lowerCAmelCase												=  val

    # finally, create HuggingFace model and load state dict
	__lowerCAmelCase												=  DetaForObjectDetection(lowerCamelCase					)
	model.load_state_dict(lowerCamelCase					)
	model.eval()

	__lowerCAmelCase												=  "cuda" if torch.cuda.is_available() else "cpu"
	model.to(lowerCamelCase					)

	# load image processor
	__lowerCAmelCase												=  DetaImageProcessor(format="coco_detection"					)

	# verify our conversion on image
	__lowerCAmelCase												=  prepare_img()
	__lowerCAmelCase												=  processor(images=lowerCamelCase				,		return_tensors="pt"					)
	__lowerCAmelCase												=  encoding["pixel_values"]
	__lowerCAmelCase												=  model(pixel_values.to(lowerCamelCase					)					)

	# verify logits
	print("Logits:"				,		outputs.logits[0, :3, :3]					)
	print("Boxes:"				,		outputs.pred_boxes[0, :3, :3]					)
	if model_name == "deta-swin-large":
		__lowerCAmelCase												=  torch.tensor(
		    [[-7.6_3_0_8, -2.8_4_8_5, -5.3_7_3_7], [-7.2_0_3_7, -4.5_5_0_5, -4.8_0_2_7], [-7.2_9_4_3, -4.2_6_1_1, -4.6_6_1_7]]					)
		__lowerCAmelCase												=  torch.tensor([[0.4_9_8_7, 0.4_9_6_9, 0.9_9_9_9], [0.2_5_4_9, 0.5_4_9_8, 0.4_8_0_5], [0.5_4_9_8, 0.2_7_5_7, 0.0_5_6_9]]					)
	elif model_name == "deta-swin-large-o365":
		__lowerCAmelCase												=  torch.tensor(
		    [[-8.0_1_2_2, -3.5_7_2_0, -4.9_7_1_7], [-8.1_5_4_7, -3.6_8_8_6, -4.6_3_8_9], [-7.6_6_1_0, -3.6_1_9_4, -5.0_1_3_4]]					)
		__lowerCAmelCase												=  torch.tensor([[0.2_5_2_3, 0.5_5_4_9, 0.4_8_8_1], [0.7_7_1_5, 0.4_1_4_9, 0.4_6_0_1], [0.5_5_0_3, 0.2_7_5_3, 0.0_5_7_5]]					)
	assert torch.allclose(outputs.logits[0, :3, :3]				,		expected_logits.to(lowerCamelCase					)				,		atol=1e-4					)
	assert torch.allclose(outputs.pred_boxes[0, :3, :3]				,		expected_boxes.to(lowerCamelCase					)				,		atol=1e-4					)
	print("Everything ok!"					)

	if pytorch_dump_folder_path:
		# Save model and processor
		logger.info(f'''Saving PyTorch model and processor to {pytorch_dump_folder_path}...'''					)
		Path(lowerCamelCase					).mkdir(exist_ok=lowerCamelCase					)
		model.save_pretrained(lowerCamelCase					)
		processor.save_pretrained(lowerCamelCase					)

	# Push to hub
	if push_to_hub:
		print("Pushing model and processor to hub..."					)
		model.push_to_hub(f'''jozhang97/{model_name}'''					)
		processor.push_to_hub(f'''jozhang97/{model_name}'''					)


if __name__ == "__main__":
	lowerCAmelCase	:		int				=   argparse.ArgumentParser()

	parser.add_argument(
	    '''--model_name''',
	    type=str,
	    default='''deta-swin-large''',
	    choices=['''deta-swin-large''', '''deta-swin-large-o365'''],
	    help='''Name of the model you\'d like to convert.''',
	)
	parser.add_argument(
	    '''--pytorch_dump_folder_path''',
	    default=None,
	    type=str,
	    help='''Path to the folder to output PyTorch model.''',
	)
	parser.add_argument(
	    '''--push_to_hub''', action='''store_true''', help='''Whether or not to push the converted model to the 🤗 hub.'''
	)
	lowerCAmelCase	:		Optional[Any]				=   parser.parse_args()
	convert_deta_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)




'''simple docstring'''



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






class 					UpperCAmelCase__	(		UpperCamelCase__ ):



					def __init__(    self     )     ->			List[str]:
						# test for the above condition
						self.test()



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

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

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

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



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



					@abstractmethod
					def        UpperCAmelCase_ (    self			,				UpperCamelCase     )     ->			Optional[int]:
						raise NotImplementedError(
						    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''     )



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



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



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



					@abstractmethod
					def        UpperCAmelCase_ (    self			,				UpperCamelCase=False     )     ->			str:
						raise NotImplementedError(
						    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''     )






class 					UpperCAmelCase__	(		UpperCamelCase__ ):



					def __init__(    self			,				UpperCamelCase     )     ->			Dict:
						super(UpperCamelCase			,				self     ).__init__()

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

						__lowerCAmelCase												=  token_ids

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



					def        UpperCAmelCase_ (    self     )     ->			Optional[int]:
						if self.completed:
							return None
						return self.token_ids[self.fulfilled_idx + 1]



					def        UpperCAmelCase_ (    self			,				UpperCamelCase     )     ->			str:
						if not isinstance(UpperCamelCase			,				UpperCamelCase     ):
							raise ValueError(F'''`token_id` has to be an `int`, but is {token_id} of type {type(UpperCamelCase     )}'''     )

						if self.completed:
							return False

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



					def        UpperCAmelCase_ (    self			,				UpperCamelCase     )     ->			int:
						if not isinstance(UpperCamelCase			,				UpperCamelCase     ):
							raise ValueError(F'''`token_id` has to be an `int`, but is {token_id} of type {type(UpperCamelCase     )}'''     )

						__lowerCAmelCase												=  False
						__lowerCAmelCase												=  False
						__lowerCAmelCase												=  False

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



					def        UpperCAmelCase_ (    self     )     ->			Union[str, Any]:
						__lowerCAmelCase												=  False
						__lowerCAmelCase												=  0



					def        UpperCAmelCase_ (    self     )     ->			Optional[int]:
						return self.seqlen - (self.fulfilled_idx + 1)



					def        UpperCAmelCase_ (    self			,				UpperCamelCase=False     )     ->			Optional[Any]:
						__lowerCAmelCase												=  PhrasalConstraint(self.token_ids     )

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

						return new_constraint






class 					UpperCAmelCase__	:



					def __init__(    self			,				UpperCamelCase			,				UpperCamelCase=True     )     ->			Optional[int]:
						__lowerCAmelCase												=  max([len(UpperCamelCase     ) for one in nested_token_ids]     )

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

								__lowerCAmelCase												=  level[token_id]

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

						__lowerCAmelCase												=  root



					def        UpperCAmelCase_ (    self			,				UpperCamelCase     )     ->			int:
						__lowerCAmelCase												=  self.trie

						for current_token in current_seq:
							__lowerCAmelCase												=  start[current_token]

						__lowerCAmelCase												=  list(start.keys()     )

						return next_tokens



					def        UpperCAmelCase_ (    self			,				UpperCamelCase     )     ->			str:
						__lowerCAmelCase												=  self.next_tokens(UpperCamelCase     )

						return len(UpperCamelCase     ) == 0



					def        UpperCAmelCase_ (    self			,				UpperCamelCase     )     ->			Optional[int]:
						__lowerCAmelCase												=  list(root.values()     )
						if len(UpperCamelCase     ) == 0:
							return 1
						else:
							return sum([self.count_leaves(UpperCamelCase     ) for nn in next_nodes]     )



					def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase     )     ->			Optional[Any]:
						__lowerCAmelCase												=  self.count_leaves(UpperCamelCase     )
						return len(UpperCamelCase     ) != leaf_count






class 					UpperCAmelCase__	(		UpperCamelCase__ ):



					def __init__(    self			,				UpperCamelCase     )     ->			List[Any]:
						super(UpperCamelCase			,				self     ).__init__()

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

						__lowerCAmelCase												=  DisjunctiveTrie(UpperCamelCase     )
						__lowerCAmelCase												=  nested_token_ids

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



					def        UpperCAmelCase_ (    self     )     ->			List[Any]:
						__lowerCAmelCase												=  self.trie.next_tokens(self.current_seq     )

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



					def        UpperCAmelCase_ (    self			,				UpperCamelCase     )     ->			List[str]:
						if not isinstance(UpperCamelCase			,				UpperCamelCase     ):
							raise ValueError(F'''`token_id` is supposed to be type `int`, but is {token_id} of type {type(UpperCamelCase     )}'''     )

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

						return token_id in next_tokens



					def        UpperCAmelCase_ (    self			,				UpperCamelCase     )     ->			int:
						if not isinstance(UpperCamelCase			,				UpperCamelCase     ):
							raise ValueError(F'''`token_id` is supposed to be type `int`, but is {token_id} of type {type(UpperCamelCase     )}'''     )

						__lowerCAmelCase												=  False
						__lowerCAmelCase												=  False
						__lowerCAmelCase												=  False

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

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

						return stepped, completed, reset



					def        UpperCAmelCase_ (    self     )     ->			Dict:
						__lowerCAmelCase												=  False
						__lowerCAmelCase												=  []



					def        UpperCAmelCase_ (    self     )     ->			int:
						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			,				UpperCamelCase=False     )     ->			Union[str, Any]:
						__lowerCAmelCase												=  DisjunctiveConstraint(self.token_ids     )

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

						return new_constraint






class 					UpperCAmelCase__	:



					def __init__(    self			,				UpperCamelCase     )     ->			Union[str, Any]:
						__lowerCAmelCase												=  constraints

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

						self.init_state()



					def        UpperCAmelCase_ (    self     )     ->			Union[str, Any]:
						__lowerCAmelCase												=  []
						__lowerCAmelCase												=  None
						__lowerCAmelCase												=  [constraint.copy(stateful=UpperCamelCase     ) for constraint in self.constraints]



					def        UpperCAmelCase_ (    self     )     ->			Optional[int]:
						__lowerCAmelCase												=  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     )     ->			List[str]:
						__lowerCAmelCase												=  []
						if self.inprogress_constraint is None:
							for constraint in self.pending_constraints:  # "pending" == "unfulfilled yet"
								__lowerCAmelCase												=  constraint.advance()
								if isinstance(UpperCamelCase			,				UpperCamelCase     ):
									token_list.append(UpperCamelCase     )
								elif isinstance(UpperCamelCase			,				UpperCamelCase     ):
									token_list.extend(UpperCamelCase     )
						else:
							__lowerCAmelCase												=  self.inprogress_constraint.advance()
							if isinstance(UpperCamelCase			,				UpperCamelCase     ):
								token_list.append(UpperCamelCase     )
							elif isinstance(UpperCamelCase			,				UpperCamelCase     ):
								token_list.extend(UpperCamelCase     )

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



					def        UpperCAmelCase_ (    self			,				UpperCamelCase     )     ->			int:
						self.init_state()

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

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



					def        UpperCAmelCase_ (    self			,				UpperCamelCase     )     ->			Dict:
						if not isinstance(UpperCamelCase			,				UpperCamelCase     ):
							raise ValueError(F'''`token_id` should be an `int`, but is `{token_id}`.'''     )

						__lowerCAmelCase							,       __lowerCAmelCase												=  False, False

						if self.completed:
							__lowerCAmelCase												=  True
							__lowerCAmelCase												=  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

							__lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase												=  self.inprogress_constraint.update(UpperCamelCase     )
							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=UpperCamelCase     )     )
								__lowerCAmelCase												=  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     )
								__lowerCAmelCase												=  None

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

									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(UpperCamelCase     )
										__lowerCAmelCase												=  None

									if not complete and stepped:
										__lowerCAmelCase												=  pending_constraint

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

										__lowerCAmelCase												=  (
										    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.

											__lowerCAmelCase												=  True

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

						return complete, stepped



					def        UpperCAmelCase_ (    self			,				UpperCamelCase=True     )     ->			str:
						__lowerCAmelCase												=  ConstraintListState(self.constraints     )  # we actually never though self.constraints objects
						# throughout this process. So it's at initialization state.

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

						return new_state