feat: support safetensors index loading (#1769)

This commit is contained in:
leejet
2026-07-10 23:23:21 +08:00
committed by GitHub
parent ead6bf521b
commit 1b042838d9
4 changed files with 84 additions and 0 deletions
+22
View File
@@ -235,6 +235,9 @@ bool ModelLoader::init_from_file(const std::string& file_path, const std::string
} else if (is_gguf_file(file_path)) {
LOG_INFO("load %s using gguf format", file_path.c_str());
return init_from_gguf_file(file_path, prefix);
} else if (ends_with(file_path, ".json")) {
LOG_INFO("load %s using safetensors index format", file_path.c_str());
return init_from_safetensors_index_file(file_path, prefix);
} else if (is_safetensors_file(file_path)) {
LOG_INFO("load %s using safetensors format", file_path.c_str());
return init_from_safetensors_file(file_path, prefix);
@@ -339,6 +342,25 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
return true;
}
bool ModelLoader::init_from_safetensors_index_file(const std::string& file_path, const std::string& prefix) {
LOG_DEBUG("init from safetensors index '%s', prefix = '%s'", file_path.c_str(), prefix.c_str());
std::vector<std::string> shard_paths;
std::string error;
if (!read_safetensors_index_file(file_path, shard_paths, &error)) {
LOG_ERROR("%s", error.c_str());
return false;
}
for (const std::string& shard_path : shard_paths) {
if (!init_from_file(shard_path, prefix)) {
return false;
}
}
return true;
}
/*================================================= TorchLegacyModelLoader ==================================================*/
bool ModelLoader::init_from_torch_legacy_file(const std::string& file_path, const std::string& prefix) {