include/tenzr/io/safetensors.h¶

// 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.

#ifndef TENZR_IO_SAFETENSORS_H_
#define TENZR_IO_SAFETENSORS_H_

#include <cstdint>
#include <memory>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

#include "tenzr/frontend/dtype.h"
#include "tenzr/frontend/shape.h"
#include "tenzr/frontend/tensor.h"

namespace tenzr::io {

  struct TensorDescriptor {
    std::string name;

    DType dtype{DType::kFloat32};

    Shape shape;

    uint64_t data_start{0};

    uint64_t data_end{0};
  };

  class SafetensorsFile {
   public:
    explicit SafetensorsFile(const std::string& filepath);

    ~SafetensorsFile();

    SafetensorsFile(const SafetensorsFile&) = delete;
    SafetensorsFile& operator=(const SafetensorsFile&) = delete;

    SafetensorsFile(SafetensorsFile&& other) noexcept;

    SafetensorsFile& operator=(SafetensorsFile&& other) noexcept;

    std::vector<std::string> TensorNames() const;

    bool HasTensor(const std::string& name) const;

    const TensorDescriptor& Descriptor(const std::string& name) const;

    const std::unordered_map<std::string, TensorDescriptor>& Descriptors() const { return descriptors_; }

    const std::unordered_map<std::string, std::string>& Metadata() const { return metadata_; }

    Tensor LoadTensor(const std::string& name) const;

    std::unordered_map<std::string, Tensor> LoadAll() const;

   private:
    void ParseHeader();

    std::string filepath_;
    int fd_{-1};
    void* mmap_data_{nullptr};
    size_t file_size_{0};
    uint64_t header_size_{0};
    const uint8_t* payload_data_{nullptr};
    std::unordered_map<std::string, TensorDescriptor> descriptors_;
    std::unordered_map<std::string, std::string> metadata_;
  };

  void SaveSafetensors(const std::string& filepath, const std::unordered_map<std::string, Tensor>& tensors,
                       const std::unordered_map<std::string, std::string>& metadata = {});

  std::unordered_map<std::string, Tensor> LoadSafetensors(const std::string& filepath);

}  // namespace tenzr::io

#endif  // TENZR_IO_SAFETENSORS_H_