// Copyright 2026 The Tenzr Project Authors
//
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "tenzr/io/safetensors.h"
#include <fcntl.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>
#include <algorithm>
#include <cctype>
#include <cstring>
#include <fstream>
#include <sstream>
#include "tenzr/common/error.h"
#include "tenzr/frontend/memory_pool.h"
#include "tenzr/io.h"
namespace tenzr::io {
static std::string DtypeToSafetensorsStr(DType dtype) {
switch (dtype) {
case DType::kBool:
return "BOOL";
case DType::kInt8:
return "I8";
case DType::kUInt8:
return "U8";
case DType::kInt16:
return "I16";
case DType::kUInt16:
return "U16";
case DType::kInt32:
return "I32";
case DType::kUInt32:
return "U32";
case DType::kInt64:
return "I64";
case DType::kUInt64:
return "U64";
case DType::kFloat16:
return "F16";
case DType::kBFloat16:
return "BF16";
case DType::kFloat32:
return "F32";
case DType::kFloat64:
return "F64";
case DType::kFloat8E4M3FN:
return "F8_E4M3";
case DType::kFloat8E5M2:
return "F8_E5M2";
case DType::kFloat8E3M4:
return "F8_E3M4";
case DType::kFloat8E4M3FNUZ:
return "F8_E4M3FNUZ";
case DType::kFloat8E4M3B11FNUZ:
return "F8_E4M3B11FNUZ";
case DType::kFloat8E5M2FNUZ:
return "F8_E5M2FNUZ";
case DType::kFloat8E8M0FNU:
return "F8_E8M0FNU";
case DType::kFloat6E2M3FN:
return "F6_E2M3";
case DType::kFloat6E3M2FN:
return "F6_E3M2";
case DType::kFloat4E2M1FN:
return "F4";
case DType::kInt1:
return "I1";
case DType::kUInt1:
return "U1";
case DType::kInt2:
return "I2";
case DType::kUInt2:
return "U2";
case DType::kInt4:
return "I4";
case DType::kUInt4:
return "U4";
case DType::kFloat8E4M3:
TENZR_THROW(
"Safetensors has no representation for float8_e4m3 (non-FN); cast to float8_e4m3fn "
"(safetensors 'F8_E4M3') before saving");
default:
TENZR_THROW("Unsupported DType for Safetensors: " + std::string(ToString(dtype)));
}
}
static DType SafetensorsStrToDtype(const std::string& str) {
if (str == "BOOL") {
return DType::kBool;
}
if (str == "I8") {
return DType::kInt8;
}
if (str == "U8") {
return DType::kUInt8;
}
if (str == "I16") {
return DType::kInt16;
}
if (str == "U16") {
return DType::kUInt16;
}
if (str == "I32") {
return DType::kInt32;
}
if (str == "U32") {
return DType::kUInt32;
}
if (str == "I64") {
return DType::kInt64;
}
if (str == "U64") {
return DType::kUInt64;
}
if (str == "F16") {
return DType::kFloat16;
}
if (str == "BF16") {
return DType::kBFloat16;
}
if (str == "F32") {
return DType::kFloat32;
}
if (str == "F64") {
return DType::kFloat64;
}
if (str == "F8_E4M3") {
return DType::kFloat8E4M3FN;
}
if (str == "F8_E5M2") {
return DType::kFloat8E5M2;
}
if (str == "F8_E3M4") {
return DType::kFloat8E3M4;
}
if (str == "F8_E4M3FNUZ") {
return DType::kFloat8E4M3FNUZ;
}
if (str == "F8_E4M3B11FNUZ") {
return DType::kFloat8E4M3B11FNUZ;
}
if (str == "F8_E5M2FNUZ") {
return DType::kFloat8E5M2FNUZ;
}
if (str == "F8_E8M0FNU") {
return DType::kFloat8E8M0FNU;
}
if (str == "F6_E2M3") {
return DType::kFloat6E2M3FN;
}
if (str == "F6_E3M2") {
return DType::kFloat6E3M2FN;
}
if (str == "F4") {
return DType::kFloat4E2M1FN;
}
if (str == "I1") {
return DType::kInt1;
}
if (str == "U1") {
return DType::kUInt1;
}
if (str == "I2") {
return DType::kInt2;
}
if (str == "U2") {
return DType::kUInt2;
}
if (str == "I4") {
return DType::kInt4;
}
if (str == "U4") {
return DType::kUInt4;
}
TENZR_THROW("Unknown Safetensors dtype string: " + str);
}
SafetensorsFile::SafetensorsFile(const std::string& filepath) : filepath_(filepath) {
fd_ = open(filepath.c_str(), O_RDONLY);
TENZR_CHECK(fd_ >= 0, "Failed to open safetensors file: " + filepath);
struct stat st;
if (fstat(fd_, &st) != 0) {
close(fd_);
TENZR_THROW("Failed to fstat safetensors file: " + filepath);
}
file_size_ = static_cast<size_t>(st.st_size);
TENZR_CHECK(file_size_ >= 8, "File too small to be a valid safetensors file: " + filepath);
mmap_data_ = mmap(nullptr, file_size_, PROT_READ, MAP_SHARED, fd_, 0);
if (mmap_data_ == MAP_FAILED) {
close(fd_);
TENZR_THROW("Failed to mmap safetensors file: " + filepath);
}
std::memcpy(&header_size_, mmap_data_, sizeof(uint64_t));
TENZR_CHECK(8 + header_size_ <= file_size_, "Corrupt safetensors file: header size (" + std::to_string(header_size_) + ") exceeds file size");
payload_data_ = static_cast<const uint8_t*>(mmap_data_) + 8 + header_size_;
ParseHeader();
}
SafetensorsFile::~SafetensorsFile() {
if (mmap_data_ && mmap_data_ != MAP_FAILED) {
munmap(mmap_data_, file_size_);
mmap_data_ = nullptr;
}
if (fd_ >= 0) {
close(fd_);
fd_ = -1;
}
}
SafetensorsFile::SafetensorsFile(SafetensorsFile&& other) noexcept
: filepath_(std::move(other.filepath_)),
fd_(other.fd_),
mmap_data_(other.mmap_data_),
file_size_(other.file_size_),
header_size_(other.header_size_),
payload_data_(other.payload_data_),
descriptors_(std::move(other.descriptors_)),
metadata_(std::move(other.metadata_)) {
other.fd_ = -1;
other.mmap_data_ = nullptr;
other.file_size_ = 0;
other.payload_data_ = nullptr;
}
SafetensorsFile& SafetensorsFile::operator=(SafetensorsFile&& other) noexcept {
if (this != &other) {
if (mmap_data_ && mmap_data_ != MAP_FAILED) {
munmap(mmap_data_, file_size_);
}
if (fd_ >= 0) {
close(fd_);
}
filepath_ = std::move(other.filepath_);
fd_ = other.fd_;
mmap_data_ = other.mmap_data_;
file_size_ = other.file_size_;
header_size_ = other.header_size_;
payload_data_ = other.payload_data_;
descriptors_ = std::move(other.descriptors_);
metadata_ = std::move(other.metadata_);
other.fd_ = -1;
other.mmap_data_ = nullptr;
other.file_size_ = 0;
other.payload_data_ = nullptr;
}
return *this;
}
void SafetensorsFile::ParseHeader() {
const char* json_ptr = static_cast<const char*>(mmap_data_) + 8;
std::string json(json_ptr, header_size_);
// Minimalist robust JSON parser for flat Safetensors header format
size_t pos = 0;
auto skip_whitespace = [&]() {
while (pos < json.size() && std::isspace(static_cast<unsigned char>(json[pos]))) {
pos++;
}
};
auto parse_string = [&]() -> std::string {
skip_whitespace();
TENZR_CHECK(pos < json.size() && json[pos] == '"', "Expected '\"' parsing string in JSON header");
pos++;
size_t start = pos;
while (pos < json.size() && json[pos] != '"') {
if (json[pos] == '\\') {
pos++;
}
pos++;
}
TENZR_CHECK(pos < json.size(), "Unterminated string in JSON header");
std::string result = json.substr(start, pos - start);
pos++; // Skip closing quote
return result;
};
auto parse_int64 = [&]() -> int64_t {
skip_whitespace();
size_t start = pos;
if (pos < json.size() && (json[pos] == '-' || json[pos] == '+')) {
pos++;
}
while (pos < json.size() && std::isdigit(json[pos])) {
pos++;
}
TENZR_CHECK(pos > start, "Expected integer in JSON header");
return std::stoll(json.substr(start, pos - start));
};
auto parse_int_array = [&]() -> std::vector<int64_t> {
skip_whitespace();
TENZR_CHECK(pos < json.size() && json[pos] == '[', "Expected '[' parsing array in JSON header");
pos++;
std::vector<int64_t> result;
skip_whitespace();
if (pos < json.size() && json[pos] == ']') {
pos++;
return result;
}
while (pos < json.size()) {
result.push_back(parse_int64());
skip_whitespace();
if (pos < json.size() && json[pos] == ',') {
pos++;
skip_whitespace();
} else if (pos < json.size() && json[pos] == ']') {
pos++;
break;
} else {
TENZR_THROW("Expected ',' or ']' in JSON array");
}
}
return result;
};
skip_whitespace();
TENZR_CHECK(pos < json.size() && json[pos] == '{', "Expected '{' at start of JSON header");
pos++;
while (pos < json.size()) {
skip_whitespace();
if (pos < json.size() && json[pos] == '}') {
pos++;
break;
}
std::string key = parse_string();
skip_whitespace();
TENZR_CHECK(pos < json.size() && json[pos] == ':', "Expected ':' after key: " + key);
pos++;
skip_whitespace();
if (key == "__metadata__") {
TENZR_CHECK(pos < json.size() && json[pos] == '{', "Expected '{' for __metadata__");
pos++;
while (pos < json.size()) {
skip_whitespace();
if (pos < json.size() && json[pos] == '}') {
pos++;
break;
}
std::string m_k = parse_string();
skip_whitespace();
TENZR_CHECK(pos < json.size() && json[pos] == ':', "Expected ':' in __metadata__");
pos++;
std::string m_v = parse_string();
metadata_[m_k] = m_v;
skip_whitespace();
if (pos < json.size() && json[pos] == ',') {
pos++;
}
}
} else {
// Tensor descriptor dictionary
TENZR_CHECK(pos < json.size() && json[pos] == '{', "Expected '{' for tensor descriptor: " + key);
pos++;
TensorDescriptor desc;
desc.name = key;
while (pos < json.size()) {
skip_whitespace();
if (pos < json.size() && json[pos] == '}') {
pos++;
break;
}
std::string field = parse_string();
skip_whitespace();
TENZR_CHECK(pos < json.size() && json[pos] == ':', "Expected ':' after field: " + field);
pos++;
skip_whitespace();
if (field == "dtype") {
std::string dtype_str = parse_string();
desc.dtype = SafetensorsStrToDtype(dtype_str);
} else if (field == "shape") {
auto shape_dims = parse_int_array();
desc.shape = Shape(shape_dims);
} else if (field == "data_offsets") {
auto offsets = parse_int_array();
TENZR_CHECK(offsets.size() == 2, "data_offsets must have exactly 2 elements");
desc.data_start = static_cast<uint64_t>(offsets[0]);
desc.data_end = static_cast<uint64_t>(offsets[1]);
} else {
// Skip unknown field
if (json[pos] == '"') {
parse_string();
} else if (json[pos] == '[') {
parse_int_array();
} else {
parse_int64();
}
}
skip_whitespace();
if (pos < json.size() && json[pos] == ',') {
pos++;
}
}
descriptors_[key] = desc;
}
skip_whitespace();
if (pos < json.size() && json[pos] == ',') {
pos++;
}
}
}
std::vector<std::string> SafetensorsFile::TensorNames() const {
std::vector<std::string> names;
for (const auto& [name, _] : descriptors_) {
names.push_back(name);
}
return names;
}
bool SafetensorsFile::HasTensor(const std::string& name) const { return descriptors_.find(name) != descriptors_.end(); }
const TensorDescriptor& SafetensorsFile::Descriptor(const std::string& name) const {
auto it = descriptors_.find(name);
TENZR_CHECK(it != descriptors_.end(), "Tensor not found in safetensors file: " + name);
return it->second;
}
Tensor SafetensorsFile::LoadTensor(const std::string& name) const {
const auto& desc = Descriptor(name);
uint64_t byte_len = desc.data_end - desc.data_start;
TENZR_CHECK(8 + header_size_ + desc.data_end <= file_size_, "Data offsets exceed file size for tensor: " + name);
const uint8_t* src_bytes = payload_data_ + desc.data_start;
auto buf = MemoryPool::Instance().Allocate(static_cast<size_t>(byte_len));
std::memcpy(buf->Data(), src_bytes, static_cast<size_t>(byte_len));
auto tensor = Tensor::Parameter(desc.shape, desc.dtype, name);
tensor.BindData(buf);
return tensor;
}
std::unordered_map<std::string, Tensor> SafetensorsFile::LoadAll() const {
std::unordered_map<std::string, Tensor> tensors;
for (const auto& [name, _] : descriptors_) {
tensors[name] = LoadTensor(name);
}
return tensors;
}
void SaveSafetensors(const std::string& filepath, const std::unordered_map<std::string, Tensor>& tensors,
const std::unordered_map<std::string, std::string>& metadata) {
std::vector<std::string> sorted_names;
for (const auto& [name, _] : tensors) {
sorted_names.push_back(name);
}
std::sort(sorted_names.begin(), sorted_names.end());
// Compute offsets
std::unordered_map<std::string, std::pair<uint64_t, uint64_t>> offsets;
uint64_t current_offset = 0;
for (const auto& name : sorted_names) {
const auto& t = tensors.at(name);
uint64_t size_bytes = static_cast<uint64_t>(t.shape().NumElements()) * SizeOf(t.Dtype());
offsets[name] = {current_offset, current_offset + size_bytes};
current_offset += size_bytes;
}
// Generate JSON header
std::ostringstream json;
json << "{";
bool first = true;
if (!metadata.empty()) {
json << "\"__metadata__\":{";
bool meta_first = true;
for (const auto& [k, v] : metadata) {
if (!meta_first) {
json << ",";
}
json << "\"" << k << "\":\"" << v << "\"";
meta_first = false;
}
json << "}";
first = false;
}
for (const auto& name : sorted_names) {
if (!first) {
json << ",";
}
first = false;
const auto& t = tensors.at(name);
const auto& off = offsets[name];
json << "\"" << name << "\":{";
json << "\"dtype\":\"" << DtypeToSafetensorsStr(t.Dtype()) << "\",";
json << "\"shape\":[";
for (size_t i = 0; i < t.shape().Rank(); ++i) {
if (i > 0) {
json << ",";
}
json << t.shape().Dim(i);
}
json << "],";
json << "\"data_offsets\":[" << off.first << "," << off.second << "]}";
}
json << "}";
std::string header_str = json.str();
uint64_t header_size = header_str.size();
std::ofstream out(filepath, std::ios::binary);
TENZR_CHECK(out.is_open(), "Failed to open destination file for safetensors: " + filepath);
out.write(reinterpret_cast<const char*>(&header_size), sizeof(header_size));
out.write(header_str.data(), static_cast<std::streamsize>(header_size));
for (const auto& name : sorted_names) {
Array t = tensors.at(name);
if (!t.HasData()) {
TENZR_CHECK(!t.IsSymbolic(), "Cannot save a symbolic tensor: " + name);
t = t.Eval();
}
TENZR_CHECK(t.HasData(), "Cannot save tensor without bound data: " + name);
auto buf = t.Data();
uint64_t size_bytes = static_cast<uint64_t>(t.shape().NumElements()) * SizeOf(t.Dtype());
out.write(reinterpret_cast<const char*>(buf->Data()), static_cast<std::streamsize>(size_bytes));
}
out.close();
}
std::unordered_map<std::string, Tensor> LoadSafetensors(const std::string& filepath) {
SafetensorsFile file(filepath);
return file.LoadAll();
}
} // namespace tenzr::io