safetensor 모델 파일 load

safetensor로 저장된 모델 불러오기

HuggingFace Trainer로 학습을 진행한 경우 기본적으로 save_safetensors 옵션이 True로 되어 있어 model.safetensors로 저장되게 된다.
이 경우 torch.load로 불러올 수 없다.

아래와 같이 불러오면 된다.

from safetensors.torch import load_model

test_model = EnsembleModel(num_labels = 5)

load_model(test_model, "/mnt/HDD8TB/PersonalityAI/FI_1head/koEnsemble_model/checkpoint-837/model.safetensors")

• Ref) https://huggingface.co/docs/safetensors/api/torch#safetensors.torch.load_model