safetensors_loader.h

Go to the documentation of this file.

#ifndef SAFETENSORS_LOADER_H
#define SAFETENSORS_LOADER_H
 
#ifdef _WIN32
#include <windows.h>
#else
#include <fcntl.h>
#include <sys/mman.h>
#include <unistd.h>
#endif
 
#include <functional>
#include <future>
#include <map>
#include <memory>
#include <mutex>
#include <nlohmann/json.hpp>
#include <queue>
#include <stdexcept>
#include <string>
#include <thread>
#include <vector>
#include <filesystem> // For directory operations, C++17
 
#include "logger.h" // Assuming Logger is accessible
 
struct ModelConfig; // Forward declaration
 
class ThreadPool; // Forward declaration
 
struct Shard {
    std::string file_path;
 
    void* mapped_data = nullptr;
 
    size_t file_size = 0;
 
    uint64_t metadata_size = 0;
 
    const uint8_t* metadata_ptr = nullptr;
 
    const uint8_t* tensor_data_block_ptr = nullptr;
 
#ifdef _WIN32
    HANDLE file_handle_ = INVALID_HANDLE_VALUE; 
    HANDLE mapping_handle_ = NULL;              
#else
    int fd_ = -1;                               
#endif
 
    explicit Shard(const std::string& fp);
 
    ~Shard();
 
    Shard(Shard&& other) noexcept;
 
    Shard& operator=(Shard&& other) noexcept;
 
    const uint8_t* get_tensor_raw_data(size_t local_offset, size_t n_bytes) const;
};
 
 
class SafeTensorsLoader {
 public:
  struct TensorInfo {
    std::string name;        
    std::string dtype;       
    std::vector<size_t> shape; 
    size_t data_offset;      
    size_t nbytes;           
    std::string shard_key;   
  };
 
  explicit SafeTensorsLoader(const std::string& model_load_path);
  
  ~SafeTensorsLoader();
 
  SafeTensorsLoader(const SafeTensorsLoader&) = delete;
  SafeTensorsLoader& operator=(const SafeTensorsLoader&) = delete;
 
  std::vector<std::string> tensor_names() const;
 
  std::vector<uint8_t> get_tensor_bytes(const std::string& name) const;
 
  const TensorInfo& get_tensor_info(const std::string& name) const;
 
  std::map<std::string, std::vector<uint8_t>> load_all_tensors_parallel() const;
 
  static bool load_model_config_from_json(const std::string& model_path_or_dir, ModelConfig& config_to_populate);
 
 private:
  std::string model_load_path_;                     
  bool is_sharded_ = false;                         
  std::map<std::string, TensorInfo> tensors_;       
  std::map<std::string, std::unique_ptr<Shard>> loaded_shards_; 
  // If sharded via an index file, this maps tensor names directly to their shard key.
  // If not sharded or sharded by pattern, this might be populated differently or less used.
  std::map<std::string, std::string> tensor_name_to_shard_key_map_; 
 
  void load_from_directory(const std::string& directory_path);
 
  void load_single_file(const std::string& file_path, const std::string& shard_key_override = "");
 
  void parse_shard_metadata(Shard& shard, const std::string& shard_key);
 
  std::vector<uint8_t> convert_tensor_data(const uint8_t* data, size_t size,
                                           const std::string& dtype) const;
 
  const Shard* get_shard_for_tensor(const std::string& tensor_name) const;
};
 
class ThreadPool {
 public:
  explicit ThreadPool(size_t num_threads);
  
  ~ThreadPool();
 
  template <class F, class... Args>
  std::future<typename std::result_of<F(Args...)>::type> submit(F&& f,
                                                               Args&&... args);
 private:
  std::vector<std::thread> workers_;           
  std::queue<std::function<void()>> tasks_;    
  std::mutex queue_mutex_;                     
  std::condition_variable condition_;          
  bool stop_ = false;                          
};
 
// Template implementation for ThreadPool::submit
template <class F, class... Args>
std::future<typename std::result_of<F(Args...)>::type> ThreadPool::submit(
    F&& f, Args&&... args) {
  using return_type = typename std::result_of<F(Args...)>::type;
 
  auto task = std::make_shared<std::packaged_task<return_type()>>(
      std::bind(std::forward<F>(f), std::forward<Args>(args)...));
 
  std::future<return_type> res = task->get_future();
  {
    std::unique_lock<std::mutex> lock(queue_mutex_);
    if (stop_) throw std::runtime_error("submit on stopped ThreadPool");
    tasks_.emplace([task]() { (*task)(); });
  }
  condition_.notify_one();
  return res;
}
 
#endif // SAFETENSORS_LOADER_H