feat: support safetensors index loading (#1769)

src/model_io/safetensors_io.cpp

@@ -3,15 +3,19 @@
 #include <algorithm>
 #include <cstdint>
 #include <exception>
+#include <filesystem>
 #include <fstream>
 #include <ostream>
 #include <string>
+#include <unordered_set>
 #include <vector>
 
 #include "binary_io.h"
 #include "core/util.h"
 #include "json.hpp"
 
+namespace fs = std::filesystem;
+
 static constexpr size_t ST_HEADER_SIZE_LEN = 8;
 
 static void set_error(std::string* error, const std::string& message) {
@@ -20,6 +24,14 @@ static void set_error(std::string* error, const std::string& message) {
     }
 }
 
+static std::string resolve_index_shard_path(const std::string& index_path, const std::string& shard_path) {
+    fs::path shard_fs_path(shard_path);
+    if (shard_fs_path.is_absolute()) {
+        return shard_fs_path.lexically_normal().string();
+    }
+    return (fs::path(index_path).parent_path() / shard_fs_path).lexically_normal().string();
+}
+
 bool is_safetensors_file(const std::string& file_path) {
     std::ifstream file(file_path, std::ios::binary);
     if (!file.is_open()) {
@@ -230,6 +242,52 @@ bool read_safetensors_file(const std::string& file_path,
     return true;
 }
 
+bool read_safetensors_index_file(const std::string& file_path,
+                                 std::vector<std::string>& shard_paths,
+                                 std::string* error) {
+    shard_paths.clear();
+
+    std::ifstream file(file_path);
+    if (!file.is_open()) {
+        set_error(error, "failed to open '" + file_path + "'");
+        return false;
+    }
+
+    nlohmann::json index;
+    try {
+        index = nlohmann::json::parse(file);
+    } catch (const std::exception&) {
+        set_error(error, "parsing safetensors index failed: '" + file_path + "'");
+        return false;
+    }
+
+    if (!index.is_object() || !index.contains("weight_map") || !index["weight_map"].is_object()) {
+        set_error(error, "invalid safetensors index '" + file_path + "'");
+        return false;
+    }
+
+    std::unordered_set<std::string> seen_shard_paths;
+    for (const auto& item : index["weight_map"].items()) {
+        if (!item.value().is_string()) {
+            set_error(error, "invalid shard path for tensor '" + item.key() + "'");
+            return false;
+        }
+
+        std::string shard_path = resolve_index_shard_path(file_path,
+                                                          item.value().get<std::string>());
+        if (seen_shard_paths.insert(shard_path).second) {
+            shard_paths.push_back(std::move(shard_path));
+        }
+    }
+
+    if (shard_paths.empty()) {
+        set_error(error, "safetensors index has no tensors: '" + file_path + "'");
+        return false;
+    }
+
+    return true;
+}
+
 static bool ggml_type_to_safetensors_dtype(ggml_type type, std::string* dtype) {
     switch (type) {
         case GGML_TYPE_F16:

src/model_io/safetensors_io.h

@@ -11,6 +11,9 @@ bool is_safetensors_file(const std::string& file_path);
 bool read_safetensors_file(const std::string& file_path,
                            std::vector<TensorStorage>& tensor_storages,
                            std::string* error = nullptr);
+bool read_safetensors_index_file(const std::string& file_path,
+                                 std::vector<std::string>& shard_paths,
+                                 std::string* error = nullptr);
 bool write_safetensors_file(const std::string& file_path,
                             const std::vector<TensorWriteInfo>& tensors,
                             std::string* error = nullptr);

src/model_loader.cpp

@@ -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) {

src/model_loader.h

@@ -43,6 +43,7 @@ protected:
 
     bool init_from_gguf_file(const std::string& file_path, const std::string& prefix = "");
     bool init_from_safetensors_file(const std::string& file_path, const std::string& prefix = "");
+    bool init_from_safetensors_index_file(const std::string& file_path, const std::string& prefix = "");
     bool init_from_torch_zip_file(const std::string& file_path, const std::string& prefix = "");
     bool init_from_torch_legacy_file(const std::string& file_path, const std::string& prefix = "");
     bool init_from_diffusers_file(const std::string& file_path, const std::string& prefix = "");