mirror of
https://github.com/curioustorvald/Terrarum-sans-bitmap.git
synced 2026-03-07 11:51:50 +09:00
270 lines
8.8 KiB
C
270 lines
8.8 KiB
C
#include "safetensor.h"
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <stdint.h>
|
|
|
|
/* Tensor registry entry */
|
|
typedef struct {
|
|
const char *name;
|
|
float *data;
|
|
int size;
|
|
int ndim;
|
|
int shape[4];
|
|
} TensorEntry;
|
|
|
|
static void collect_tensors(Network *net, TensorEntry *entries, int *count) {
|
|
int n = 0;
|
|
#define ADD(nm, layer, field) do { \
|
|
entries[n].name = nm; \
|
|
entries[n].data = net->layer.field->data; \
|
|
entries[n].size = net->layer.field->size; \
|
|
entries[n].ndim = net->layer.field->ndim; \
|
|
for (int i = 0; i < net->layer.field->ndim; i++) \
|
|
entries[n].shape[i] = net->layer.field->shape[i]; \
|
|
n++; \
|
|
} while(0)
|
|
|
|
ADD("conv1.weight", conv1, weight);
|
|
ADD("conv1.bias", conv1, bias);
|
|
ADD("conv2.weight", conv2, weight);
|
|
ADD("conv2.bias", conv2, bias);
|
|
ADD("fc1.weight", fc1, weight);
|
|
ADD("fc1.bias", fc1, bias);
|
|
ADD("head_shape.weight", head_shape, weight);
|
|
ADD("head_shape.bias", head_shape, bias);
|
|
ADD("head_ytype.weight", head_ytype, weight);
|
|
ADD("head_ytype.bias", head_ytype, bias);
|
|
ADD("head_lowheight.weight", head_lowheight, weight);
|
|
ADD("head_lowheight.bias", head_lowheight, bias);
|
|
#undef ADD
|
|
*count = n;
|
|
}
|
|
|
|
int safetensor_save(const char *path, Network *net, int total_samples, int epochs, float val_loss) {
|
|
TensorEntry entries[12];
|
|
int count;
|
|
collect_tensors(net, entries, &count);
|
|
|
|
/* Build JSON header */
|
|
char header[8192];
|
|
int pos = 0;
|
|
pos += snprintf(header + pos, sizeof(header) - (size_t)pos, "{");
|
|
|
|
/* metadata */
|
|
pos += snprintf(header + pos, sizeof(header) - (size_t)pos,
|
|
"\"__metadata__\":{\"samples\":\"%d\",\"epochs\":\"%d\",\"val_loss\":\"%.6f\"},",
|
|
total_samples, epochs, (double)val_loss);
|
|
|
|
/* tensor entries */
|
|
size_t data_offset = 0;
|
|
for (int i = 0; i < count; i++) {
|
|
size_t byte_size = (size_t)entries[i].size * sizeof(float);
|
|
|
|
pos += snprintf(header + pos, sizeof(header) - (size_t)pos,
|
|
"\"%s\":{\"dtype\":\"F32\",\"shape\":[", entries[i].name);
|
|
|
|
for (int d = 0; d < entries[i].ndim; d++) {
|
|
if (d > 0) pos += snprintf(header + pos, sizeof(header) - (size_t)pos, ",");
|
|
pos += snprintf(header + pos, sizeof(header) - (size_t)pos, "%d", entries[i].shape[d]);
|
|
}
|
|
|
|
pos += snprintf(header + pos, sizeof(header) - (size_t)pos,
|
|
"],\"data_offsets\":[%zu,%zu]}", data_offset, data_offset + byte_size);
|
|
|
|
if (i < count - 1)
|
|
pos += snprintf(header + pos, sizeof(header) - (size_t)pos, ",");
|
|
|
|
data_offset += byte_size;
|
|
}
|
|
|
|
pos += snprintf(header + pos, sizeof(header) - (size_t)pos, "}");
|
|
|
|
/* Pad header to 8-byte alignment */
|
|
size_t header_len = (size_t)pos;
|
|
size_t padded = (header_len + 7) & ~(size_t)7;
|
|
while (header_len < padded) {
|
|
header[header_len++] = ' ';
|
|
}
|
|
|
|
FILE *f = fopen(path, "wb");
|
|
if (!f) {
|
|
fprintf(stderr, "Error: cannot open %s for writing\n", path);
|
|
return -1;
|
|
}
|
|
|
|
/* 8-byte LE header length */
|
|
uint64_t hlen = (uint64_t)header_len;
|
|
fwrite(&hlen, 8, 1, f);
|
|
|
|
/* JSON header */
|
|
fwrite(header, 1, header_len, f);
|
|
|
|
/* Raw tensor data */
|
|
for (int i = 0; i < count; i++) {
|
|
fwrite(entries[i].data, sizeof(float), (size_t)entries[i].size, f);
|
|
}
|
|
|
|
fclose(f);
|
|
printf("Saved model to %s (%zu bytes)\n", path, 8 + header_len + data_offset);
|
|
return 0;
|
|
}
|
|
|
|
/* Minimal JSON parser: find tensor by name, extract data_offsets */
|
|
static int find_tensor_offsets(const char *json, size_t json_len, const char *name,
|
|
size_t *off_start, size_t *off_end) {
|
|
/* Search for "name": */
|
|
size_t nlen = strlen(name);
|
|
for (size_t i = 0; i + nlen + 3 < json_len; i++) {
|
|
if (json[i] == '"' && strncmp(json + i + 1, name, nlen) == 0 && json[i + 1 + nlen] == '"') {
|
|
/* Found the key, now find data_offsets */
|
|
const char *doff = strstr(json + i, "\"data_offsets\"");
|
|
if (!doff || (size_t)(doff - json) > json_len) return -1;
|
|
const char *bracket = strchr(doff, '[');
|
|
if (!bracket) return -1;
|
|
if (sscanf(bracket, "[%zu,%zu]", off_start, off_end) != 2) return -1;
|
|
return 0;
|
|
}
|
|
}
|
|
return -1;
|
|
}
|
|
|
|
int safetensor_load(const char *path, Network *net) {
|
|
FILE *f = fopen(path, "rb");
|
|
if (!f) {
|
|
fprintf(stderr, "Error: cannot open %s\n", path);
|
|
return -1;
|
|
}
|
|
|
|
uint64_t header_len;
|
|
if (fread(&header_len, 8, 1, f) != 1) { fclose(f); return -1; }
|
|
|
|
char *json = malloc((size_t)header_len + 1);
|
|
if (fread(json, 1, (size_t)header_len, f) != (size_t)header_len) {
|
|
free(json);
|
|
fclose(f);
|
|
return -1;
|
|
}
|
|
json[header_len] = '\0';
|
|
|
|
long data_start = 8 + (long)header_len;
|
|
|
|
TensorEntry entries[12];
|
|
int count;
|
|
collect_tensors(net, entries, &count);
|
|
|
|
for (int i = 0; i < count; i++) {
|
|
size_t off_start, off_end;
|
|
if (find_tensor_offsets(json, (size_t)header_len, entries[i].name, &off_start, &off_end) != 0) {
|
|
fprintf(stderr, "Error: tensor '%s' not found in %s\n", entries[i].name, path);
|
|
free(json);
|
|
fclose(f);
|
|
return -1;
|
|
}
|
|
|
|
size_t byte_size = off_end - off_start;
|
|
if (byte_size != (size_t)entries[i].size * sizeof(float)) {
|
|
fprintf(stderr, "Error: size mismatch for '%s': expected %zu, got %zu\n",
|
|
entries[i].name, (size_t)entries[i].size * sizeof(float), byte_size);
|
|
free(json);
|
|
fclose(f);
|
|
return -1;
|
|
}
|
|
|
|
fseek(f, data_start + (long)off_start, SEEK_SET);
|
|
if (fread(entries[i].data, 1, byte_size, f) != byte_size) {
|
|
fprintf(stderr, "Error: failed to read tensor '%s'\n", entries[i].name);
|
|
free(json);
|
|
fclose(f);
|
|
return -1;
|
|
}
|
|
}
|
|
|
|
free(json);
|
|
fclose(f);
|
|
return 0;
|
|
}
|
|
|
|
int safetensor_stats(const char *path) {
|
|
FILE *f = fopen(path, "rb");
|
|
if (!f) {
|
|
fprintf(stderr, "Error: cannot open %s\n", path);
|
|
return -1;
|
|
}
|
|
|
|
uint64_t header_len;
|
|
if (fread(&header_len, 8, 1, f) != 1) { fclose(f); return -1; }
|
|
|
|
char *json = malloc((size_t)header_len + 1);
|
|
if (fread(json, 1, (size_t)header_len, f) != (size_t)header_len) {
|
|
free(json);
|
|
fclose(f);
|
|
return -1;
|
|
}
|
|
json[header_len] = '\0';
|
|
fclose(f);
|
|
|
|
printf("Model: %s\n", path);
|
|
printf("Header length: %lu bytes\n", (unsigned long)header_len);
|
|
|
|
/* Extract a JSON string value: find "key":"value" and return value */
|
|
/* Helper: find value for key within metadata block */
|
|
const char *meta = strstr(json, "\"__metadata__\"");
|
|
if (meta) {
|
|
const char *keys[] = {"samples", "epochs", "val_loss"};
|
|
const char *labels[] = {"Training samples", "Epochs", "Validation loss"};
|
|
for (int k = 0; k < 3; k++) {
|
|
char search[64];
|
|
snprintf(search, sizeof(search), "\"%s\"", keys[k]);
|
|
const char *found = strstr(meta, search);
|
|
if (!found) continue;
|
|
/* skip past key and colon to opening quote of value */
|
|
const char *colon = strchr(found + strlen(search), ':');
|
|
if (!colon) continue;
|
|
const char *vstart = strchr(colon, '"');
|
|
if (!vstart) continue;
|
|
vstart++;
|
|
const char *vend = strchr(vstart, '"');
|
|
if (!vend) continue;
|
|
printf("%s: %.*s\n", labels[k], (int)(vend - vstart), vstart);
|
|
}
|
|
}
|
|
|
|
/* List tensors */
|
|
const char *tensor_names[] = {
|
|
"conv1.weight", "conv1.bias", "conv2.weight", "conv2.bias",
|
|
"fc1.weight", "fc1.bias",
|
|
"head_shape.weight", "head_shape.bias",
|
|
"head_ytype.weight", "head_ytype.bias",
|
|
"head_lowheight.weight", "head_lowheight.bias"
|
|
};
|
|
|
|
int total_params = 0;
|
|
printf("\nTensors:\n");
|
|
for (int i = 0; i < 12; i++) {
|
|
size_t off_start, off_end;
|
|
if (find_tensor_offsets(json, (size_t)header_len, tensor_names[i], &off_start, &off_end) == 0) {
|
|
int params = (int)(off_end - off_start) / 4;
|
|
total_params += params;
|
|
|
|
/* Extract shape */
|
|
const char *key = strstr(json, tensor_names[i]);
|
|
if (key) {
|
|
const char *shp = strstr(key, "\"shape\"");
|
|
if (shp) {
|
|
const char *br = strchr(shp, '[');
|
|
const char *bre = strchr(shp, ']');
|
|
if (br && bre) {
|
|
printf(" %-28s shape=[%.*s] params=%d\n",
|
|
tensor_names[i], (int)(bre - br - 1), br + 1, params);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
printf("\nTotal parameters: %d (%.1f KB as float32)\n", total_params, (float)total_params * 4.0f / 1024.0f);
|
|
|
|
free(json);
|
|
return 0;
|
|
}
|