Create upload_model.py

upload_model.py

@@ -0,0 +1,39 @@
+from transformers import AutoConfig, AutoModel
+import torch
+
+from .modeling_chada_vit import ChAdaViTModel
+from .config_chada_vit import ChAdaViTConfig
+
+AutoConfig.register("chadavit", ChAdaViTConfig)
+AutoModel.register(ChAdaViTConfig, ChAdaViTModel)
+
+ChAdaViTConfig.register_for_auto_class()
+ChAdaViTModel.register_for_auto_class("AutoModel")
+
+config = ChAdaViTConfig()
+
+
+def adjust_keys(state_dict):
+    new_state_dict = {}
+    for key in state_dict:
+        new_key = key
+        if "encoder" in key:
+            new_key = new_key.replace("encoder", "backbone")
+        if "backbone" in key:
+            new_key = new_key.replace("backbone.", "")
+        new_state_dict[new_key] = state_dict[key]
+    return new_state_dict
+
+
+# Initialize model
+model = AutoModel.from_config(config)
+
+# Load state dictionary
+ckpt_path = "CKPT_PTH.ckpt"
+state_dict = torch.load(ckpt_path, map_location="cpu")["state_dict"]
+
+# Adjust state dictionary keys and load it into the model
+adjusted_state_dict = adjust_keys(state_dict)
+model.load_state_dict(adjusted_state_dict, strict=False)
+
+model.push_to_hub("chadavit16-moyen")