"""
    #################################################################################################################################################################################################
    # This code library is an adaptation of the original Transformers and was designed, developed and programmed by Sapiens Technology®.                                                            #
    # Any alteration and/or disclosure of this code without prior authorization is strictly prohibited and is subject to legal action that will be forwarded by the Sapiens Technology® legal team. #
    # This set of algorithms aims to download, train, fine-tune and/or infer large language models from various sources and slopes.                                                                 #
    #################################################################################################################################################################################################
"""
import argparse
import re
from pathlib import Path
import requests
import torch
from PIL import Image
from torchvision.transforms import CenterCrop, Compose, Normalize, Resize, ToTensor
from sapiens_transformers import (EfficientFormerConfig, EfficientFormerForImageClassificationWithTeacher, EfficientFormerImageProcessor)
from sapiens_transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, PILImageResampling
def rename_key(old_name, num_meta4D_last_stage):
    new_name = old_name
    if "patch_embed" in old_name:
        _, layer, param = old_name.split(".")
        if layer == "0": new_name = old_name.replace("0", "convolution1")
        elif layer == "1": new_name = old_name.replace("1", "batchnorm_before")
        elif layer == "3": new_name = old_name.replace("3", "convolution2")
        else: new_name = old_name.replace("4", "batchnorm_after")
    if "network" in old_name and re.search(r"\d\.\d", old_name):
        two_digit_num = r"\b\d{2}\b"
        if bool(re.search(two_digit_num, old_name)): match = re.search(r"\d\.\d\d.", old_name).group()
        else: match = re.search(r"\d\.\d.", old_name).group()
        if int(match[0]) < 6:
            trimmed_name = old_name.replace(match, "")
            trimmed_name = trimmed_name.replace("network", match[0] + ".meta4D_layers.blocks." + match[2:-1])
            new_name = "intermediate_stages." + trimmed_name
        else:
            trimmed_name = old_name.replace(match, "")
            if int(match[2]) < num_meta4D_last_stage: trimmed_name = trimmed_name.replace("network", "meta4D_layers.blocks." + match[2])
            else:
                layer_index = str(int(match[2]) - num_meta4D_last_stage)
                trimmed_name = trimmed_name.replace("network", "meta3D_layers.blocks." + layer_index)
                if "norm1" in old_name: trimmed_name = trimmed_name.replace("norm1", "layernorm1")
                elif "norm2" in old_name: trimmed_name = trimmed_name.replace("norm2", "layernorm2")
                elif "fc1" in old_name: trimmed_name = trimmed_name.replace("fc1", "linear_in")
                elif "fc2" in old_name: trimmed_name = trimmed_name.replace("fc2", "linear_out")
            new_name = "last_stage." + trimmed_name
    elif "network" in old_name and re.search(r".\d.", old_name): new_name = old_name.replace("network", "intermediate_stages")
    if "fc" in new_name: new_name = new_name.replace("fc", "convolution")
    elif ("norm1" in new_name) and ("layernorm1" not in new_name): new_name = new_name.replace("norm1", "batchnorm_before")
    elif ("norm2" in new_name) and ("layernorm2" not in new_name): new_name = new_name.replace("norm2", "batchnorm_after")
    if "proj" in new_name: new_name = new_name.replace("proj", "projection")
    if "dist_head" in new_name: new_name = new_name.replace("dist_head", "distillation_classifier")
    elif "head" in new_name: new_name = new_name.replace("head", "classifier")
    elif "patch_embed" in new_name: new_name = "efficientformer." + new_name
    elif new_name == "norm.weight" or new_name == "norm.bias":
        new_name = new_name.replace("norm", "layernorm")
        new_name = "efficientformer." + new_name
    else: new_name = "efficientformer.encoder." + new_name
    return new_name
def convert_torch_checkpoint(checkpoint, num_meta4D_last_stage):
    for key in checkpoint.copy().keys():
        val = checkpoint.pop(key)
        checkpoint[rename_key(key, num_meta4D_last_stage)] = val
    return checkpoint
def prepare_img():
    url = "http://images.cocodataset.org/val2017/000000039769.jpg"
    image = Image.open(requests.get(url, stream=True).raw)
    return image
def convert_efficientformer_checkpoint(checkpoint_path: Path, efficientformer_config_file: Path, pytorch_dump_path: Path, push_to_hub: bool):
    orig_state_dict = torch.load(checkpoint_path, map_location="cpu")["model"]
    config = EfficientFormerConfig.from_json_file(efficientformer_config_file)
    model = EfficientFormerForImageClassificationWithTeacher(config)
    model_name = "_".join(checkpoint_path.split("/")[-1].split(".")[0].split("_")[:-1])
    num_meta4D_last_stage = config.depths[-1] - config.num_meta3d_blocks + 1
    new_state_dict = convert_torch_checkpoint(orig_state_dict, num_meta4D_last_stage)
    model.load_state_dict(new_state_dict)
    model.eval()
    pillow_resamplings = {"bilinear": PILImageResampling.BILINEAR, "bicubic": PILImageResampling.BICUBIC, "nearest": PILImageResampling.NEAREST}
    image = prepare_img()
    image_size = 256
    crop_size = 224
    processor = EfficientFormerImageProcessor(size={"shortest_edge": image_size}, crop_size={"height": crop_size, "width": crop_size}, resample=pillow_resamplings["bicubic"])
    pixel_values = processor(images=image, return_tensors="pt").pixel_values
    image_transforms = Compose([Resize(image_size, interpolation=pillow_resamplings["bicubic"]), CenterCrop(crop_size), ToTensor(), Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)])
    original_pixel_values = image_transforms(image).unsqueeze(0)
    assert torch.allclose(original_pixel_values, pixel_values)
    outputs = model(pixel_values)
    logits = outputs.logits
    expected_shape = (1, 1000)
    if "l1" in model_name:
        expected_logits = torch.Tensor([-0.1312, 0.4353, -1.0499, -0.5124, 0.4183, -0.6793, -1.3777, -0.0893, -0.7358, -2.4328])
        assert torch.allclose(logits[0, :10], expected_logits, atol=1e-3)
        assert logits.shape == expected_shape
    elif "l3" in model_name:
        expected_logits = torch.Tensor([-1.3150, -1.5456, -1.2556, -0.8496, -0.7127, -0.7897, -0.9728, -0.3052, 0.3751, -0.3127])
        assert torch.allclose(logits[0, :10], expected_logits, atol=1e-3)
        assert logits.shape == expected_shape
    elif "l7" in model_name:
        expected_logits = torch.Tensor([-1.0283, -1.4131, -0.5644, -1.3115, -0.5785, -1.2049, -0.7528, 0.1992, -0.3822, -0.0878])
        assert logits.shape == expected_shape
    else: raise ValueError(f"Unknown model checkpoint: {checkpoint_path}. Supported version of efficientformer are l1, l3 and l7")
    Path(pytorch_dump_path).mkdir(exist_ok=True)
    model.save_pretrained(pytorch_dump_path)
    print(f"Checkpoint successfuly converted. Model saved at {pytorch_dump_path}")
    processor.save_pretrained(pytorch_dump_path)
    print(f"Processor successfuly saved at {pytorch_dump_path}")
    if push_to_hub:
        print("Pushing model to the hub...")
        model.push_to_hub(repo_id=f"Bearnardd/{pytorch_dump_path}", commit_message="Add model", use_temp_dir=True)
        processor.push_to_hub(repo_id=f"Bearnardd/{pytorch_dump_path}", commit_message="Add image processor", use_temp_dir=True)
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--pytorch_model_path", default=None, type=str, required=True, help="Path to EfficientFormer pytorch checkpoint.")
    parser.add_argument("--config_file", default=None, type=str, required=True, help="The json file for EfficientFormer model config.")
    parser.add_argument("--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model.")
    parser.add_argument("--push_to_hub", action="store_true", help="Push model and image processor to the hub")
    parser.add_argument("--no-push_to_hub", dest="push_to_hub", action="store_false", help="Do not push model and image processor to the hub")
    parser.set_defaults(push_to_hub=True)
    args = parser.parse_args()
    convert_efficientformer_checkpoint(checkpoint_path=args.pytorch_model_path, efficientformer_config_file=args.config_file, pytorch_dump_path=args.pytorch_dump_path, push_to_hub=args.push_to_hub)
"""
    #################################################################################################################################################################################################
    # This code library is an adaptation of the original Transformers and was designed, developed and programmed by Sapiens Technology®.                                                            #
    # Any alteration and/or disclosure of this code without prior authorization is strictly prohibited and is subject to legal action that will be forwarded by the Sapiens Technology® legal team. #
    # This set of algorithms aims to download, train, fine-tune and/or infer large language models from various sources and slopes.                                                                 #
    #################################################################################################################################################################################################
"""
