




import argparse
import re

import torch
from CLAP import create_model

from transformers import AutoFeatureExtractor, ClapConfig, ClapModel


UpperCamelCase__					    =    {
    """text_branch""": """text_model""",
    """audio_branch""": """audio_model.audio_encoder""",
    """attn""": """attention.self""",
    """self.proj""": """output.dense""",
    """attention.self_mask""": """attn_mask""",
    """mlp.fc1""": """intermediate.dense""",
    """mlp.fc2""": """output.dense""",
    """norm1""": """layernorm_before""",
    """norm2""": """layernorm_after""",
    """bn0""": """batch_norm""",
}

UpperCamelCase__					    =    AutoFeatureExtractor.from_pretrained("""laion/clap-htsat-unfused""", truncation="""rand_trunc""")







def 			_a   (		SCREAMING_SNAKE_CASE_							:       int   ,   SCREAMING_SNAKE_CASE_							:       Any=False  ):
					__lowerCAmelCase     ,							__lowerCAmelCase            = create_model(
					    "HTSAT-tiny"   ,   "roberta"   ,   SCREAMING_SNAKE_CASE_   ,   precision="fp32"   ,   device="cuda:0" if torch.cuda.is_available() else "cpu"   ,   enable_fusion=SCREAMING_SNAKE_CASE_   ,   fusion_type="aff_2d" if enable_fusion else None   ,   )
					return model, model_cfg







def 			_a   (		SCREAMING_SNAKE_CASE_							:       List[Any]  ):
					__lowerCAmelCase            = {}

					__lowerCAmelCase            = R".*sequential.(\d+).*"
					__lowerCAmelCase            = R".*_projection.(\d+).*"

					for key, value in state_dict.items():
										# check if any key needs to be modified
										for key_to_modify, new_key in KEYS_TO_MODIFY_MAPPING.items():
															if key_to_modify in key:
																				__lowerCAmelCase            = key.replace(SCREAMING_SNAKE_CASE_   ,   SCREAMING_SNAKE_CASE_  )

										if re.match(SCREAMING_SNAKE_CASE_   ,   SCREAMING_SNAKE_CASE_  ):
															# replace sequential layers with list
															__lowerCAmelCase            = re.match(SCREAMING_SNAKE_CASE_   ,   SCREAMING_SNAKE_CASE_  ).group(1  )

															__lowerCAmelCase            = key.replace(F"""sequential.{sequential_layer}."""   ,   F"""layers.{int(SCREAMING_SNAKE_CASE_  )//3}.linear."""  )
										elif re.match(SCREAMING_SNAKE_CASE_   ,   SCREAMING_SNAKE_CASE_  ):
															__lowerCAmelCase            = int(re.match(SCREAMING_SNAKE_CASE_   ,   SCREAMING_SNAKE_CASE_  ).group(1  )  )

															# Because in CLAP they use `nn.Sequential`...
															__lowerCAmelCase            = 1 if projecton_layer == 0 else 2

															__lowerCAmelCase            = key.replace(F"""_projection.{projecton_layer}."""   ,   F"""_projection.linear{transformers_projection_layer}."""  )

										if "audio" and "qkv" in key:
															# split qkv into query key and value
															__lowerCAmelCase            = value
															__lowerCAmelCase            = mixed_qkv.size(0  ) // 3

															__lowerCAmelCase            = mixed_qkv[:qkv_dim]
															__lowerCAmelCase            = mixed_qkv[qkv_dim : qkv_dim * 2]
															__lowerCAmelCase            = mixed_qkv[qkv_dim * 2 :]

															__lowerCAmelCase            = query_layer
															__lowerCAmelCase            = key_layer
															__lowerCAmelCase            = value_layer
										else:
															__lowerCAmelCase            = value

					return model_state_dict







def 			_a   (		SCREAMING_SNAKE_CASE_							:       Union[str, Any]   ,   SCREAMING_SNAKE_CASE_							:       Optional[int]   ,   SCREAMING_SNAKE_CASE_							:       Optional[Any]   ,   SCREAMING_SNAKE_CASE_							:       Any=False  ):
					__lowerCAmelCase     ,							__lowerCAmelCase            = init_clap(SCREAMING_SNAKE_CASE_   ,   enable_fusion=SCREAMING_SNAKE_CASE_  )

					clap_model.eval()
					__lowerCAmelCase            = clap_model.state_dict()
					__lowerCAmelCase            = rename_state_dict(SCREAMING_SNAKE_CASE_  )

					__lowerCAmelCase            = ClapConfig()
					__lowerCAmelCase            = enable_fusion
					__lowerCAmelCase            = ClapModel(SCREAMING_SNAKE_CASE_  )

					# ignore the spectrogram embedding layer
					model.load_state_dict(SCREAMING_SNAKE_CASE_   ,   strict=SCREAMING_SNAKE_CASE_  )

					model.save_pretrained(SCREAMING_SNAKE_CASE_  )
					transformers_config.save_pretrained(SCREAMING_SNAKE_CASE_  )


if __name__ == "__main__":
					UpperCamelCase__					    =    argparse.ArgumentParser()
					parser.add_argument("""--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model.""")
					parser.add_argument("""--checkpoint_path""", default=None, type=str, help="""Path to fairseq checkpoint""")
					parser.add_argument("""--config_path""", default=None, type=str, help="""Path to hf config.json of model to convert""")
					parser.add_argument("""--enable_fusion""", action="""store_true""", help="""Whether to enable fusion or not""")
					UpperCamelCase__					    =    parser.parse_args()

					convert_clap_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.enable_fusion)










from math import ceil, sqrt







def 			_a   (		SCREAMING_SNAKE_CASE_							:       int = 1_00_00_00  ):
					__lowerCAmelCase            = 0

					for outer_width in range(3   ,   (limit // 4) + 2  ):
										if outer_width**2 > limit:
															__lowerCAmelCase            = max(ceil(sqrt(outer_width**2 - limit  )  )   ,   1  )
										else:
															__lowerCAmelCase            = 1
										if (outer_width - hole_width_lower_bound) % 2:
															hole_width_lower_bound += 1

										answer += (outer_width - hole_width_lower_bound - 2) // 2 + 1

					return answer


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




