




import argparse
from pathlib import Path

import torch

from transformers import OPTConfig, OPTModel
from transformers.utils import logging


logging.set_verbosity_info()
__UpperCAmelCase     							=						logging.get_logger(__name__)




def 					lowercase__				(					__snake_case   :							int  ):


					'''simple docstring'''


					UpperCAmelCase_   :				List[Any]						=						torch.load(__snake_case					,	map_location='cpu'  )
					if "model" in sd.keys():
										UpperCAmelCase_   :				str						=						torch.load(__snake_case					,	map_location='cpu'  )['model']

					# pop unnecessary weights
					UpperCAmelCase_   :				Optional[int]						=						[
					    'decoder.version',
					    'decoder.output_projection.weight',
					]
					for key in keys_to_delete:
										if key in sd:
															sd.pop(__snake_case  )

					UpperCAmelCase_   :				int						=						{
					    'decoder.project_in_dim.weight': 'decoder.project_in.weight',
					    'decoder.project_out_dim.weight': 'decoder.project_out.weight',
					    'decoder.layer_norm.weight': 'decoder.final_layer_norm.weight',
					    'decoder.layer_norm.bias': 'decoder.final_layer_norm.bias',
					}
					for old_key, new_key in keys_to_rename.items():
										if old_key in sd:
															UpperCAmelCase_   :				List[str]						=						sd.pop(__snake_case  )

					UpperCAmelCase_   :				Optional[int]						=						list(sd.keys()  )
					for key in keys:
										if ".qkv_proj." in key:
															UpperCAmelCase_   :				Dict						=						sd[key]
															# We split QKV in separate Q,K,V

															UpperCAmelCase_   :				List[Any]						=						key.replace('.qkv_proj.'					,	'.q_proj.'  )
															UpperCAmelCase_   :				Dict						=						key.replace('.qkv_proj.'					,	'.k_proj.'  )
															UpperCAmelCase_   :				Union[str, Any]						=						key.replace('.qkv_proj.'					,	'.v_proj.'  )

															UpperCAmelCase_   :				Union[str, Any]						=						value.shape[0]
															assert depth % 3 == 0
															# `SequeuceParallelTransformerBlock` has QKV weight is separated in K,V,Q despite the naming:
															# https://cs.github.com/facebookresearch/metaseq/blob/51871bd73cd04c038f239ea2a26db1d7f6b37927/metaseq/modules/sequence_parallel_transformer_layer.py#L97
															UpperCAmelCase_   :				str						=						torch.split(__snake_case					,	depth // 3					,	dim=0  )

															UpperCAmelCase_   :				Optional[Any]						=						q
															UpperCAmelCase_   :				List[Any]						=						k
															UpperCAmelCase_   :				Any						=						v
															del sd[key]

					return sd




@torch.no_grad()
def 					lowercase__				(					__snake_case   :							Dict					,	__snake_case   :							Any					,	__snake_case   :							Any=None  ):


					'''simple docstring'''


					UpperCAmelCase_   :				Optional[int]						=						load_checkpoint(__snake_case  )

					if config is not None:
										UpperCAmelCase_   :				str						=						OPTConfig.from_pretrained(__snake_case  )
					else:
										UpperCAmelCase_   :				Optional[int]						=						OPTConfig()

					UpperCAmelCase_   :				Tuple						=						OPTModel(__snake_case  ).half().eval()
					model.load_state_dict(__snake_case  )

					# Check results
					Path(__snake_case  ).mkdir(exist_ok=__snake_case  )
					model.save_pretrained(__snake_case  )


if __name__ == "__main__":
		__UpperCAmelCase     							=						argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    '--fairseq_path',
		    type=str,
		    help=(
		        'path to fairseq checkpoint in correct format. You can find all checkpoints in the correct format here:'
		        ' https://huggingface.co/models?other=opt_metasq'
		    ),
		)
		parser.add_argument('--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model.')
		parser.add_argument('--hf_config', default=None, type=str, help='Define HF config.')
		__UpperCAmelCase     							=						parser.parse_args()
		convert_opt_checkpoint(args.fairseq_path, args.pytorch_dump_folder_path, config=args.hf_config)














from collections import defaultdict
from math import ceil, sqrt




def 					lowercase__				(					__snake_case   :							int = 1_000_000					,	__snake_case   :							int = 10  ):


					'''simple docstring'''


					UpperCAmelCase_   :				defaultdict						=						defaultdict(__snake_case  )

					for outer_width in range(3					,	(t_limit // 4) + 2  ):
										if outer_width * outer_width > t_limit:
															UpperCAmelCase_   :				Union[str, Any]						=						max(
															    ceil(sqrt(outer_width * outer_width - t_limit  )  )					,	1  )
										else:
															UpperCAmelCase_   :				int						=						1

										hole_width_lower_bound += (outer_width - hole_width_lower_bound) % 2

										for hole_width in range(__snake_case					,	outer_width - 1					,	2  ):
															count[outer_width * outer_width - hole_width * hole_width] += 1

					return sum(1 for n in count.values() if 1 <= n <= 10  )


if __name__ == "__main__":
		print(F'{solution() = }')






