code

app.py

@@ -6,6 +6,32 @@ import os
 
 os.system('huggingface-cli download matteogeniaccio/phi-4 --local-dir ./phi-4 --include "phi-4/*"')
 
+from safetensors.torch import load_file, save_file
+
+@spaces.GPU
+def merge_safetensors(input_dir, output_file):
+    # 获取所有分片文件
+    files = sorted([f for f in os.listdir(input_dir) if f.startswith('model-') and f.endswith('.safetensors')])
+    
+    # 合并所有张量
+    merged_state_dict = {}
+    for file in files:
+        file_path = os.path.join(input_dir, file)
+        print(f"Loading {file}...")
+        state_dict = load_file(file_path)
+        merged_state_dict.update(state_dict)
+    
+    # 保存合并后的文件
+    print(f"Saving merged model to {output_file}...")
+    save_file(merged_state_dict, output_file)
+    print("Done!")
+
+# 使用示例
+input_dir = "./phi-4/phi-4"  # 包含分片文件的目录
+output_file = "./phi-4/phi-4/model.safetensors"  # 合并后的文件路径
+
+merge_safetensors(input_dir, output_file)
+
 # 加载 phi-4 模型和 tokenizer
 torch.random.manual_seed(0)