158 lines
4.6 KiB
C
158 lines
4.6 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
#ifndef FOUNDRY_SAFETENSORS_H
|
|
#define FOUNDRY_SAFETENSORS_H
|
|
|
|
#include <stdint.h>
|
|
#include <stdbool.h>
|
|
#include <stddef.h>
|
|
|
|
#ifdef __cplusplus
|
|
extern "C" {
|
|
#endif
|
|
|
|
/*
|
|
* Safe tensor serialization format.
|
|
* Replaces HuggingFace safetensors.
|
|
*
|
|
* Format:
|
|
* - 8 bytes: header length (little-endian u64)
|
|
* - N bytes: JSON header with tensor metadata
|
|
* - Remaining: raw tensor data
|
|
*/
|
|
|
|
/* ============================================================= */
|
|
/* Types */
|
|
/* ============================================================= */
|
|
|
|
/* Safetensors data types (matches safetensors spec) */
|
|
typedef enum {
|
|
FS_DTYPE_F32,
|
|
FS_DTYPE_F64,
|
|
FS_DTYPE_F16,
|
|
FS_DTYPE_BF16,
|
|
FS_DTYPE_I8,
|
|
FS_DTYPE_I16,
|
|
FS_DTYPE_I32,
|
|
FS_DTYPE_I64,
|
|
FS_DTYPE_U8,
|
|
FS_DTYPE_U16,
|
|
FS_DTYPE_U32,
|
|
FS_DTYPE_U64,
|
|
FS_DTYPE_BOOL,
|
|
FS_DTYPE_UNKNOWN
|
|
} fs_dtype;
|
|
|
|
/* Tensor info */
|
|
typedef struct {
|
|
char* name;
|
|
fs_dtype dtype;
|
|
int64_t* shape;
|
|
int ndim;
|
|
size_t data_offset;
|
|
size_t data_size;
|
|
} fs_tensor_info;
|
|
|
|
/* Safetensors file handle */
|
|
typedef struct fs_file fs_file;
|
|
|
|
/* Writer for creating new safetensors files */
|
|
typedef struct fs_writer fs_writer;
|
|
|
|
/* ============================================================= */
|
|
/* Reading */
|
|
/* ============================================================= */
|
|
|
|
/* Open safetensors file (memory-maps the file) */
|
|
fs_file* fs_open(const char* path);
|
|
|
|
/* Open from memory buffer (does not copy) */
|
|
fs_file* fs_open_buffer(const void* data, size_t size);
|
|
|
|
/* Close file and free resources */
|
|
void fs_close(fs_file* f);
|
|
|
|
/* Get number of tensors in file */
|
|
int fs_num_tensors(fs_file* f);
|
|
|
|
/* Get list of tensor names (returns array of strings, caller must free array but not strings) */
|
|
const char** fs_tensor_names(fs_file* f, int* count);
|
|
|
|
/* Check if tensor exists */
|
|
bool fs_has_tensor(fs_file* f, const char* name);
|
|
|
|
/* Get tensor info (returns NULL if not found) */
|
|
fs_tensor_info* fs_get_info(fs_file* f, const char* name);
|
|
|
|
/* Free tensor info */
|
|
void fs_free_info(fs_tensor_info* info);
|
|
|
|
/* Get raw tensor data pointer (zero-copy, points into mmap'd region) */
|
|
const void* fs_get_data(fs_file* f, const char* name);
|
|
|
|
/* Get tensor data with type checking (returns NULL if dtype mismatch) */
|
|
const float* fs_get_f32(fs_file* f, const char* name);
|
|
const double* fs_get_f64(fs_file* f, const char* name);
|
|
const int32_t* fs_get_i32(fs_file* f, const char* name);
|
|
const int64_t* fs_get_i64(fs_file* f, const char* name);
|
|
const uint8_t* fs_get_u8(fs_file* f, const char* name);
|
|
|
|
/* ============================================================= */
|
|
/* Metadata */
|
|
/* ============================================================= */
|
|
|
|
/* Get metadata value (returns NULL if not found) */
|
|
const char* fs_get_metadata(fs_file* f, const char* key);
|
|
|
|
/* Get all metadata keys */
|
|
const char** fs_metadata_keys(fs_file* f, int* count);
|
|
|
|
/* ============================================================= */
|
|
/* Writing */
|
|
/* ============================================================= */
|
|
|
|
/* Create new safetensors writer */
|
|
fs_writer* fs_writer_new(void);
|
|
|
|
/* Free writer */
|
|
void fs_writer_free(fs_writer* w);
|
|
|
|
/* Add metadata */
|
|
void fs_writer_add_metadata(fs_writer* w, const char* key, const char* value);
|
|
|
|
/* Add tensor to writer */
|
|
void fs_writer_add_tensor(fs_writer* w, const char* name, fs_dtype dtype,
|
|
const int64_t* shape, int ndim, const void* data);
|
|
|
|
/* Add float32 tensor (convenience) */
|
|
void fs_writer_add_f32(fs_writer* w, const char* name,
|
|
const int64_t* shape, int ndim, const float* data);
|
|
|
|
/* Write to file */
|
|
bool fs_writer_save(fs_writer* w, const char* path);
|
|
|
|
/* Write to memory buffer (caller must free) */
|
|
void* fs_writer_to_buffer(fs_writer* w, size_t* size);
|
|
|
|
/* ============================================================= */
|
|
/* Utility */
|
|
/* ============================================================= */
|
|
|
|
/* Get dtype name string */
|
|
const char* fs_dtype_name(fs_dtype dtype);
|
|
|
|
/* Get dtype size in bytes */
|
|
size_t fs_dtype_size(fs_dtype dtype);
|
|
|
|
/* Parse dtype from string */
|
|
fs_dtype fs_dtype_from_string(const char* s);
|
|
|
|
/* Print file summary */
|
|
void fs_print(fs_file* f);
|
|
|
|
#ifdef __cplusplus
|
|
}
|
|
#endif
|
|
|
|
#endif /* FOUNDRY_SAFETENSORS_H */
|