

from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from ..utils import BaseOutput, is_torch_version, randn_tensor
from .attention_processor import SpatialNorm
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block






@dataclass
class _UpperCamelCase				(				__A ):
			'''simple docstring'''







			lowerCamelCase__							     =42






class _UpperCamelCase				(				nn.Module ):
			'''simple docstring'''







			def __init__(      self				:			int   ,   a				:			List[Any]=3   ,   a				:			Optional[int]=3   ,   a				:			Dict=("DownEncoderBlock2D",)   ,   a				:			Any=(64,)   ,   a				:			List[Any]=2   ,   a				:			Dict=32   ,   a				:			str="silu"   ,   a				:			Optional[Any]=True   ,   )				->							int:

				"""simple docstring"""

				super().__init__()
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							layers_per_block

				SCREAMING_SNAKE_CASE :   Tuple					  =							torch.nn.Convad(
				    a   ,   block_out_channels[0]   ,   kernel_size=3   ,   stride=1   ,   padding=1   ,   )

				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							None
				SCREAMING_SNAKE_CASE :   Any					  =							nn.ModuleList([]						)

				# down
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							block_out_channels[0]
				for i, down_block_type in enumerate(a						):
					SCREAMING_SNAKE_CASE :   str					  =							output_channel
					SCREAMING_SNAKE_CASE :   int					  =							block_out_channels[i]
					SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							i == len(a						) - 1

					SCREAMING_SNAKE_CASE :   Optional[int]					  =							get_down_block(
					    a   ,   num_layers=self.layers_per_block   ,   in_channels=a   ,   out_channels=a   ,   add_downsample=not is_final_block   ,   resnet_eps=1e-6   ,   downsample_padding=0   ,   resnet_act_fn=a   ,   resnet_groups=a   ,   attention_head_dim=a   ,   temb_channels=a   ,   )
					self.down_blocks.append(a						)

				# mid
				SCREAMING_SNAKE_CASE :   str					  =							UNetMidBlockaD(
				    in_channels=block_out_channels[-1]   ,   resnet_eps=1e-6   ,   resnet_act_fn=a   ,   output_scale_factor=1   ,   resnet_time_scale_shift="default"   ,   attention_head_dim=block_out_channels[-1]   ,   resnet_groups=a   ,   temb_channels=a   ,   )

				# out
				SCREAMING_SNAKE_CASE :   str					  =							nn.GroupNorm(num_channels=block_out_channels[-1]   ,   num_groups=a   ,   eps=1e-6						)
				SCREAMING_SNAKE_CASE :   Dict					  =							nn.SiLU()

				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							2 * out_channels if double_z else out_channels
				SCREAMING_SNAKE_CASE :   Tuple					  =							nn.Convad(block_out_channels[-1]   ,   a   ,   3   ,   padding=1						)

				SCREAMING_SNAKE_CASE :   Any					  =							False






			def 				__UpperCamelCase			(      self				:			List[Any]   ,   a				:			str						)				->							Any:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   List[str]					  =							x
				SCREAMING_SNAKE_CASE :   List[str]					  =							self.conv_in(a						)

				if self.training and self.gradient_checkpointing:

					def create_custom_forward(a				:			Any						):
						def custom_forward(*a				:			Optional[int]						):
							return module(*a						)

						return custom_forward

					# down
					if is_torch_version(">="   ,   "1.11.0"						):
						for down_block in self.down_blocks:
							SCREAMING_SNAKE_CASE :   str					  =							torch.utils.checkpoint.checkpoint(
							    create_custom_forward(a						)   ,   a   ,   use_reentrant=a						)
						# middle
						SCREAMING_SNAKE_CASE :   Tuple					  =							torch.utils.checkpoint.checkpoint(
						    create_custom_forward(self.mid_block						)   ,   a   ,   use_reentrant=a						)
					else:
						for down_block in self.down_blocks:
							SCREAMING_SNAKE_CASE :   Any					  =							torch.utils.checkpoint.checkpoint(create_custom_forward(a						)   ,   a						)
						# middle
						SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block						)   ,   a						)

				else:
					# down
					for down_block in self.down_blocks:
						SCREAMING_SNAKE_CASE :   List[Any]					  =							down_block(a						)

					# middle
					SCREAMING_SNAKE_CASE :   Optional[Any]					  =							self.mid_block(a						)

				# post-process
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							self.conv_norm_out(a						)
				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							self.conv_act(a						)
				SCREAMING_SNAKE_CASE :   str					  =							self.conv_out(a						)

				return sample






class _UpperCamelCase				(				nn.Module ):
			'''simple docstring'''







			def __init__(      self				:			List[Any]   ,   a				:			Dict=3   ,   a				:			int=3   ,   a				:			Optional[Any]=("UpDecoderBlock2D",)   ,   a				:			Tuple=(64,)   ,   a				:			int=2   ,   a				:			Any=32   ,   a				:			Optional[int]="silu"   ,   a				:			int="group"   ,   )				->							int:

				"""simple docstring"""

				super().__init__()
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							layers_per_block

				SCREAMING_SNAKE_CASE :   Any					  =							nn.Convad(
				    a   ,   block_out_channels[-1]   ,   kernel_size=3   ,   stride=1   ,   padding=1   ,   )

				SCREAMING_SNAKE_CASE :   List[Any]					  =							None
				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							nn.ModuleList([]						)

				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							in_channels if norm_type == "spatial" else None

				# mid
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							UNetMidBlockaD(
				    in_channels=block_out_channels[-1]   ,   resnet_eps=1e-6   ,   resnet_act_fn=a   ,   output_scale_factor=1   ,   resnet_time_scale_shift="default" if norm_type == "group" else norm_type   ,   attention_head_dim=block_out_channels[-1]   ,   resnet_groups=a   ,   temb_channels=a   ,   )

				# up
				SCREAMING_SNAKE_CASE :   Any					  =							list(reversed(a						)						)
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							reversed_block_out_channels[0]
				for i, up_block_type in enumerate(a						):
					SCREAMING_SNAKE_CASE :   Tuple					  =							output_channel
					SCREAMING_SNAKE_CASE :   List[str]					  =							reversed_block_out_channels[i]

					SCREAMING_SNAKE_CASE :   int					  =							i == len(a						) - 1

					SCREAMING_SNAKE_CASE :   List[str]					  =							get_up_block(
					    a   ,   num_layers=self.layers_per_block + 1   ,   in_channels=a   ,   out_channels=a   ,   prev_output_channel=a   ,   add_upsample=not is_final_block   ,   resnet_eps=1e-6   ,   resnet_act_fn=a   ,   resnet_groups=a   ,   attention_head_dim=a   ,   temb_channels=a   ,   resnet_time_scale_shift=a   ,   )
					self.up_blocks.append(a						)
					SCREAMING_SNAKE_CASE :   Any					  =							output_channel

				# out
				if norm_type == "spatial":
					SCREAMING_SNAKE_CASE :   Optional[int]					  =							SpatialNorm(block_out_channels[0]   ,   a						)
				else:
					SCREAMING_SNAKE_CASE :   List[str]					  =							nn.GroupNorm(num_channels=block_out_channels[0]   ,   num_groups=a   ,   eps=1e-6						)
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							nn.SiLU()
				SCREAMING_SNAKE_CASE :   Any					  =							nn.Convad(block_out_channels[0]   ,   a   ,   3   ,   padding=1						)

				SCREAMING_SNAKE_CASE :   int					  =							False






			def 				__UpperCamelCase			(      self				:			Tuple   ,   a				:			Union[str, Any]   ,   a				:			List[Any]=None						)				->							Optional[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   List[Any]					  =							z
				SCREAMING_SNAKE_CASE :   int					  =							self.conv_in(a						)

				SCREAMING_SNAKE_CASE :   str					  =							next(iter(self.up_blocks.parameters()						)						).dtype
				if self.training and self.gradient_checkpointing:

					def create_custom_forward(a				:			Dict						):
						def custom_forward(*a				:			List[Any]						):
							return module(*a						)

						return custom_forward

					if is_torch_version(">="   ,   "1.11.0"						):
						# middle
						SCREAMING_SNAKE_CASE :   str					  =							torch.utils.checkpoint.checkpoint(
						    create_custom_forward(self.mid_block						)   ,   a   ,   a   ,   use_reentrant=a						)
						SCREAMING_SNAKE_CASE :   List[Any]					  =							sample.to(a						)

						# up
						for up_block in self.up_blocks:
							SCREAMING_SNAKE_CASE :   Optional[int]					  =							torch.utils.checkpoint.checkpoint(
							    create_custom_forward(a						)   ,   a   ,   a   ,   use_reentrant=a						)
					else:
						# middle
						SCREAMING_SNAKE_CASE :   Any					  =							torch.utils.checkpoint.checkpoint(
						    create_custom_forward(self.mid_block						)   ,   a   ,   a						)
						SCREAMING_SNAKE_CASE :   Dict					  =							sample.to(a						)

						# up
						for up_block in self.up_blocks:
							SCREAMING_SNAKE_CASE :   int					  =							torch.utils.checkpoint.checkpoint(create_custom_forward(a						)   ,   a   ,   a						)
				else:
					# middle
					SCREAMING_SNAKE_CASE :   Optional[int]					  =							self.mid_block(a   ,   a						)
					SCREAMING_SNAKE_CASE :   List[Any]					  =							sample.to(a						)

					# up
					for up_block in self.up_blocks:
						SCREAMING_SNAKE_CASE :   Dict					  =							up_block(a   ,   a						)

        # post-process
				if latent_embeds is None:
					SCREAMING_SNAKE_CASE :   Optional[int]					  =							self.conv_norm_out(a						)
				else:
					SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							self.conv_norm_out(a   ,   a						)
				SCREAMING_SNAKE_CASE :   List[Any]					  =							self.conv_act(a						)
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							self.conv_out(a						)

				return sample






class _UpperCamelCase				(				nn.Module ):
			'''simple docstring'''







			def __init__(      self				:			Tuple   ,   a				:			str   ,   a				:			Optional[int]   ,   a				:			Optional[int]   ,   a				:			Tuple=None   ,   a				:			Optional[Any]="random"   ,   a				:			Any=False   ,   a				:			List[str]=True						)				->							Union[str, Any]:

				"""simple docstring"""

				super().__init__()
				SCREAMING_SNAKE_CASE :   List[Any]					  =							n_e
				SCREAMING_SNAKE_CASE :   str					  =							vq_embed_dim
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							beta
				SCREAMING_SNAKE_CASE :   Any					  =							legacy

				SCREAMING_SNAKE_CASE :   List[Any]					  =							nn.Embedding(self.n_e   ,   self.vq_embed_dim						)
				self.embedding.weight.data.uniform_(-1.0 / self.n_e   ,   1.0 / self.n_e						)

				SCREAMING_SNAKE_CASE :   List[Any]					  =							remap
				if self.remap is not None:
					self.register_buffer("used"   ,   torch.tensor(np.load(self.remap						)						)						)
					SCREAMING_SNAKE_CASE :   Dict					  =							self.used.shape[0]
					SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							unknown_index  # "random" or "extra" or integer
					if self.unknown_index == "extra":
						SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							self.re_embed
						SCREAMING_SNAKE_CASE :   List[Any]					  =							self.re_embed + 1
					print(
					    F"Remapping {self.n_e} indices to {self.re_embed} indices. "
					    F"Using {self.unknown_index} for unknown indices."						)
				else:
					SCREAMING_SNAKE_CASE :   List[str]					  =							n_e

				SCREAMING_SNAKE_CASE :   Any					  =							sane_index_shape
			def 				__UpperCamelCase			(      self				:			Tuple   ,   a				:			Any						)				->							List[str]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							inds.shape
				assert len(a						) > 1
				SCREAMING_SNAKE_CASE :   Dict					  =							inds.reshape(ishape[0]   ,   -1						)
				SCREAMING_SNAKE_CASE :   List[str]					  =							self.used.to(a						)
				SCREAMING_SNAKE_CASE :   List[str]					  =							(inds[:, :, None] == used[None, None, ...]).long()
				SCREAMING_SNAKE_CASE :   Tuple					  =							match.argmax(-1						)
				SCREAMING_SNAKE_CASE :   Dict					  =							match.sum(2						) < 1
				if self.unknown_index == "random":
					SCREAMING_SNAKE_CASE :   Tuple					  =							torch.randint(0   ,   self.re_embed   ,   size=new[unknown].shape						).to(device=new.device						)
				else:
					SCREAMING_SNAKE_CASE :   Optional[int]					  =							self.unknown_index
				return new.reshape(a						)
			def 				__UpperCamelCase			(      self				:			str   ,   a				:			List[Any]						)				->							Dict:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Optional[int]					  =							inds.shape
				assert len(a						) > 1
				SCREAMING_SNAKE_CASE :   List[str]					  =							inds.reshape(ishape[0]   ,   -1						)
				SCREAMING_SNAKE_CASE :   List[str]					  =							self.used.to(a						)
				if self.re_embed > self.used.shape[0]:  # extra token
					SCREAMING_SNAKE_CASE :   Dict					  =							0  # simply set to zero
				SCREAMING_SNAKE_CASE :   Any					  =							torch.gather(used[None, :][inds.shape[0] * [0], :]   ,   1   ,   a						)
				return back.reshape(a						)
			def 				__UpperCamelCase			(      self				:			Union[str, Any]   ,   a				:			Tuple						)				->							Any:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Any					  =							z.permute(0   ,   2   ,   3   ,   1						).contiguous()
				SCREAMING_SNAKE_CASE :   str					  =							z.view(-1   ,   self.vq_embed_dim						)

				# distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
				SCREAMING_SNAKE_CASE :   str					  =							torch.argmin(torch.cdist(a   ,   self.embedding.weight						)   ,   dim=1						)

				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							self.embedding(a						).view(z.shape						)
				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							None
				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							None

				# compute loss for embedding
				if not self.legacy:
					SCREAMING_SNAKE_CASE :   Optional[Any]					  =							self.beta * torch.mean((z_q.detach() - z) ** 2						) + torch.mean((z_q - z.detach()) ** 2						)
				else:
					SCREAMING_SNAKE_CASE :   Any					  =							torch.mean((z_q.detach() - z) ** 2						) + self.beta * torch.mean((z_q - z.detach()) ** 2						)

				# preserve gradients
				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							z + (z_q - z).detach()

				# reshape back to match original input shape
				SCREAMING_SNAKE_CASE :   str					  =							z_q.permute(0   ,   3   ,   1   ,   2						).contiguous()

				if self.remap is not None:
					SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							min_encoding_indices.reshape(z.shape[0]   ,   -1						)  # add batch axis
					SCREAMING_SNAKE_CASE :   Optional[Any]					  =							self.remap_to_used(a						)
					SCREAMING_SNAKE_CASE :   Optional[int]					  =							min_encoding_indices.reshape(-1   ,   1						)  # flatten

				if self.sane_index_shape:
					SCREAMING_SNAKE_CASE :   int					  =							min_encoding_indices.reshape(z_q.shape[0]   ,   z_q.shape[2]   ,   z_q.shape[3]						)

				return z_q, loss, (perplexity, min_encodings, min_encoding_indices)






			def 				__UpperCamelCase			(      self				:			Any   ,   a				:			Dict   ,   a				:			str						)				->							List[str]:

				"""simple docstring"""

				if self.remap is not None:
					SCREAMING_SNAKE_CASE :   Any					  =							indices.reshape(shape[0]   ,   -1						)  # add batch axis
					SCREAMING_SNAKE_CASE :   int					  =							self.unmap_to_all(a						)
					SCREAMING_SNAKE_CASE :   str					  =							indices.reshape(-1						)  # flatten again

				# get quantized latent vectors
				SCREAMING_SNAKE_CASE :   List[Any]					  =							self.embedding(a						)

				if shape is not None:
					SCREAMING_SNAKE_CASE :   Any					  =							z_q.view(a						)
					# reshape back to match original input shape
					SCREAMING_SNAKE_CASE :   int					  =							z_q.permute(0   ,   3   ,   1   ,   2						).contiguous()

				return z_q






class _UpperCamelCase				(				__A ):
			'''simple docstring'''







			def __init__(      self				:			List[str]   ,   a				:			List[str]   ,   a				:			Dict=False						)				->							int:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							parameters
				SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE :   Any					  =							torch.chunk(a   ,   2   ,   dim=1						)
				SCREAMING_SNAKE_CASE :   Any					  =							torch.clamp(self.logvar   ,   -30.0   ,   20.0						)
				SCREAMING_SNAKE_CASE :   str					  =							deterministic
				SCREAMING_SNAKE_CASE :   int					  =							torch.exp(0.5 * self.logvar						)
				SCREAMING_SNAKE_CASE :   Tuple					  =							torch.exp(self.logvar						)
				if self.deterministic:
					SCREAMING_SNAKE_CASE :   Tuple					  =							torch.zeros_like(
					    self.mean   ,   device=self.parameters.device   ,   dtype=self.parameters.dtype						)
			def 				__UpperCamelCase			(      self				:			Optional[int]   ,   a				:			Optional[torch.Generator] = None						)				->							torch.FloatTensor:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Any					  =							randn_tensor(
				    self.mean.shape   ,   generator=a   ,   device=self.parameters.device   ,   dtype=self.parameters.dtype						)
				SCREAMING_SNAKE_CASE :   int					  =							self.mean + self.std * sample
				return x
			def 				__UpperCamelCase			(      self				:			Optional[Any]   ,   a				:			str=None						)				->							Any:

				"""simple docstring"""

				if self.deterministic:
					return torch.Tensor([0.0]						)
				else:
					if other is None:
						return 0.5 * torch.sum(torch.pow(self.mean   ,   2						) + self.var - 1.0 - self.logvar   ,   dim=[1, 2, 3]						)
					else:
						return 0.5 * torch.sum(
						    torch.pow(self.mean - other.mean   ,   2						) / other.var
						    + self.var / other.var
						    - 1.0
						    - self.logvar
						    + other.logvar   ,   dim=[1, 2, 3]   ,   )
			def 				__UpperCamelCase			(      self				:			List[Any]   ,   a				:			Union[str, Any]   ,   a				:			Dict=[1, 2, 3]						)				->							int:

				"""simple docstring"""

				if self.deterministic:
					return torch.Tensor([0.0]						)
				SCREAMING_SNAKE_CASE :   List[str]					  =							np.log(2.0 * np.pi						)
				return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean   ,   2						) / self.var   ,   dim=a						)






			def 				__UpperCamelCase			(      self				:			Optional[Any]						)				->							Dict:

				"""simple docstring"""

				return self.mean


import argparse
import logging
from collections import namedtuple

import torch
from model_bertabs import BertAbsSummarizer
from models.model_builder import AbsSummarizer  # The authors' implementation

from transformers import BertTokenizer


logging.basicConfig(level=logging.INFO)
a_      =     logging.getLogger(__name__)


a_      =     'Hello world! cécé herlolip'


a_      =     namedtuple(
    'BertAbsConfig',
    [
        'temp_dir',
        'large',
        'use_bert_emb',
        'finetune_bert',
        'encoder',
        'share_emb',
        'max_pos',
        'enc_layers',
        'enc_hidden_size',
        'enc_heads',
        'enc_ff_size',
        'enc_dropout',
        'dec_layers',
        'dec_hidden_size',
        'dec_heads',
        'dec_ff_size',
        'dec_dropout',
    ],
)

def 			lowerCamelCase__			(       _a				,      _a):
	SCREAMING_SNAKE_CASE :   List[Any]					  =							BertAbsConfig(
	    temp_dir="."				,      finetune_bert=_a				,      large=_a				,      share_emb=_a				,      use_bert_emb=_a				,      encoder="bert"				,      max_pos=512				,      enc_layers=6				,      enc_hidden_size=512				,      enc_heads=8				,      enc_ff_size=512				,      enc_dropout=0.2				,      dec_layers=6				,      dec_hidden_size=768				,      dec_heads=8				,      dec_ff_size=2048				,      dec_dropout=0.2				,      )
	SCREAMING_SNAKE_CASE :   Dict					  =							torch.load(_a				,      lambda _a				,      _a: storage)
	SCREAMING_SNAKE_CASE :   str					  =							AbsSummarizer(_a				,      torch.device("cpu")				,      _a)
	original.eval()

	SCREAMING_SNAKE_CASE :   List[str]					  =							BertAbsSummarizer(_a				,      torch.device("cpu"))
	new_model.eval()

	# -------------------
	# Convert the weights
	# -------------------

	logging.info("convert the model")
	new_model.bert.load_state_dict(original.bert.state_dict())
	new_model.decoder.load_state_dict(original.decoder.state_dict())
	new_model.generator.load_state_dict(original.generator.state_dict())

	# ----------------------------------
	# Make sure the outpus are identical
	# ----------------------------------

	logging.info("Make sure that the models' outputs are identical")
	SCREAMING_SNAKE_CASE :   List[str]					  =							BertTokenizer.from_pretrained("bert-base-uncased")

	# prepare the model inputs
	SCREAMING_SNAKE_CASE :   List[str]					  =							tokenizer.encode("This is sample éàalj'-.")
	encoder_input_ids.extend([tokenizer.pad_token_id] * (512 - len(_a)))
	SCREAMING_SNAKE_CASE :   int					  =							torch.tensor(_a).unsqueeze(0)
	SCREAMING_SNAKE_CASE :   List[Any]					  =							tokenizer.encode("This is sample 3 éàalj'-.")
	decoder_input_ids.extend([tokenizer.pad_token_id] * (512 - len(_a)))
	SCREAMING_SNAKE_CASE :   int					  =							torch.tensor(_a).unsqueeze(0)

	# failsafe to make sure the weights reset does not affect the
	# loaded weights.
	assert torch.max(torch.abs(original.generator[0].weight - new_model.generator[0].weight)) == 0

	# forward pass
	SCREAMING_SNAKE_CASE :   List[Any]					  =							encoder_input_ids
	SCREAMING_SNAKE_CASE :   List[Any]					  =							decoder_input_ids
	SCREAMING_SNAKE_CASE :   Dict					  =							None
	SCREAMING_SNAKE_CASE :   Optional[Any]					  =							None
	SCREAMING_SNAKE_CASE :   List[Any]					  =							None
	SCREAMING_SNAKE_CASE :   Optional[int]					  =							None
	SCREAMING_SNAKE_CASE :   Dict					  =							None

	# The original model does not apply the geneator layer immediatly but rather in
	# the beam search (where it combines softmax + linear layer). Since we already
	# apply the softmax in our generation process we only apply the linear layer here.
	# We make sure that the outputs of the full stack are identical
	SCREAMING_SNAKE_CASE :   Optional[int]					  =							original(_a				,      _a				,      _a				,      _a				,      _a				,      _a				,      _a)[0]
	SCREAMING_SNAKE_CASE :   Dict					  =							original.generator(_a)

	SCREAMING_SNAKE_CASE :   Any					  =							new_model(
	    _a				,      _a				,      _a				,      _a				,      _a)[0]
	SCREAMING_SNAKE_CASE :   Tuple					  =							new_model.generator(_a)

	SCREAMING_SNAKE_CASE :   List[Any]					  =							torch.max(torch.abs(output_converted_model - output_original_model)).item()
	print("Maximum absolute difference beween weights: {:.2f}".format(_a))
	SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							torch.max(torch.abs(output_converted_generator - output_original_generator)).item()
	print("Maximum absolute difference beween weights: {:.2f}".format(_a))

	SCREAMING_SNAKE_CASE :   int					  =							torch.allclose(_a				,      _a				,      atol=1E-3)
	if are_identical:
		logging.info("all weights are equal up to 1e-3")
	else:
		raise ValueError("the weights are different. The new model is likely different from the original one.")

	# The model has been saved with torch.save(model) and this is bound to the exact
	# directory structure. We save the state_dict instead.
	logging.info("saving the model's state dictionary")
	torch.save(
	    new_model.state_dict()				,      "./bertabs-finetuned-cnndm-extractive-abstractive-summarization/pytorch_model.bin")


if __name__ == "__main__":
			a_      =     argparse.ArgumentParser()
			parser.add_argument(
			    '--bertabs_checkpoint_path',
			    default=None,
			    type=str,
			    required=True,
			    help='Path the official PyTorch dump.',
			)
			parser.add_argument(
			    '--pytorch_dump_folder_path',
			    default=None,
			    type=str,
			    required=True,
			    help='Path to the output PyTorch model.',
			)
			a_      =     parser.parse_args()

			convert_bertabs_checkpoints(
			    args.bertabs_checkpoint_path,
			    args.pytorch_dump_folder_path,
			)