pth to st

scripts/convert_pth_to_safetensors.py

@@ -1,18 +1,11 @@
 import os
 
 import torch
-from transformers import AutoModel
+from safetensors.torch import save_file
+
 
 checkpoint_dir = '../out/pretrain-core-3/hf'
 output_dir = '../out/pretrain-core-3/hf'
 
-# Load model
 state_dict = torch.load(os.path.join(checkpoint_dir, 'model.pth'))
-
-model = AutoModel.from_pretrained(
-    checkpoint_dir,
-    state_dict=state_dict,
-)
-
-# Save .safetensors files
-model.save_pretrained(output_dir)
+save_file(state_dict, os.path.join(output_dir, 'model.safetensors'))