from huggingface_hub import notebook_login
from transformers import AutoConfig, AutoModel
from autoencoder_model.modeling_autoencoder import AutoEncoder, AutoEncoderConfig

notebook_login()

# Register Huggingface Model
AutoEncoderConfig.register_for_auto_class()
AutoEncoder.register_for_auto_class("AutoModel")

AutoConfig.register("autoencoder", AutoEncoderConfig)
AutoModel.register(AutoEncoderConfig, AutoModel)

# Create Model
autoencoder = AutoEncoder(AutoEncoderConfig())
autoencoder.push_to_hub("autoencoder")

# Download Model
config = AutoConfig.from_pretrained("amaye15/autoencoder", trust_remote_code = True)
autoencoder = AutoModel.from_config(config, trust_remote_code = True)