[Huggingface] Safetensors Basic Usage

UNGGI LEE·2023년 10월 6일
0

먼저, Huggingface에서 모델을 불러옵니다.

from transformers import AutoTokenizer, AutoModelForCausalLM

model_path = "EleutherAI/polyglot-ko-5.8b"
model = AutoModelForCausalLM.from_pretrained(model_path)

다음으로, safetensors를 활용하여 불러온 모델을 safetensors로 저장합니다.

from safetensors.torch import load_model, save_model

# save_model(저장할 모델, 저장할 경로)
save_model(model, "my_model.safetensors")

그럼 모델이 원하는 위치에 safetensors 형태로 저장되게 됩니다.

reference

https://huggingface.co/docs/safetensors/index

profile
NLP Researcher & ML Engineer @i-Scream Edu

0개의 댓글