safetensors.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <cstddef>
#include <cstdint>
#include <map>
#include <string>
#include <utility>
#include <vector>
 
#include "pulsatrix/tensor.hpp"
 
namespace pulsatrix {
 
enum class SafetensorsDtype { Bool, U8, I8, I16, U16, I32, U32, I64, U64, F8_E4M3, F8_E5M2, F16, BF16, F32, F64 };
 
struct SafetensorsTensorInfo {
    SafetensorsDtype dtype;
    std::vector<int64_t> shape;
    uint64_t data_begin;
    uint64_t data_end;
};
 
class SafetensorsFile {
public:
    [[nodiscard]] static SafetensorsFile Read(const std::string& path);
 
    [[nodiscard]] static SafetensorsFile Parse(std::vector<uint8_t> bytes);
 
    [[nodiscard]] const std::vector<std::string>& names() const { return names_; }
 
    [[nodiscard]] bool contains(const std::string& name) const { return infos_.count(name) != 0; }
 
    [[nodiscard]] const SafetensorsTensorInfo& info(const std::string& name) const;
 
    [[nodiscard]] const std::map<std::string, std::string>& metadata() const { return metadata_; }
 
    [[nodiscard]] std::pair<const uint8_t*, size_t> bytes(const std::string& name) const;
 
    [[nodiscard]] Tensor tensor(const std::string& name, DeviceBackend* backend) const;
 
private:
    SafetensorsFile() = default;
 
    std::vector<uint8_t> bytes_;
    size_t data_start_ = 0;
    std::vector<std::string> names_;
    std::map<std::string, SafetensorsTensorInfo> infos_;
    std::map<std::string, std::string> metadata_;
};
 
[[nodiscard]] std::vector<uint8_t> SerializeSafetensors(const std::vector<std::pair<std::string, const Tensor*>>& tensors,
                                                        const std::map<std::string, std::string>& metadata = {});
 
void WriteSafetensors(const std::string& path, const std::vector<std::pair<std::string, const Tensor*>>& tensors,
                      const std::map<std::string, std::string>& metadata = {});
 
}  // namespace pulsatrix