Fix handle_safetensor

handle_safetensors.py

@@ -3,11 +3,10 @@ from safetensors.torch import safe_open
 import os
 import torch
 import argparse
+import json
 from transformers import AutoModelForCausalLM
 
 def save_model_at_once(model, save_dir):
-    import pdb
-    pdb.set_trace()
     tensors = {k:v for k, v in model.state_dict().items()}
     path = os.path.join(save_dir, "model.safetensors")
     save_file(tensors, path)
@@ -47,12 +46,20 @@ if __name__ == "__main__":
     args = parser.parse_args()
     
     model = AutoModelForCausalLM.from_pretrained(args.model_path)
+    
     print("Model loaded")
 
     if not os.path.exists(args.save_dir):
         from pathlib import Path
         Path(args.save_dir).mkdir(parents=True, exist_ok=True)
 
+    conf = dict(sorted(model.config.to_diff_dict().items(), key=lambda x: x[0]))
+    del conf['architectures']
+    del conf['model_type']
+    conf['torch_dtype'] = "bfloat16"
+    with open(os.path.join(args.save_dir, "config.json"), "w") as f:
+        json.dump(conf, f, indent=2)
+
     load_path = args.save_dir
     if args.n_file == 1:
         save_model_at_once(model, args.save_dir)