ex_pth.c

Go to the documentation of this file.

/*
 * Nilorea Library
 * Copyright (C) 2005-2026 Castagnier Mickael
 *
 * 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
 *
 *     http://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.
 *
 * SPDX-License-Identifier: Apache-2.0
 */

 
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
 
#include <zlib.h>
 
#include "nilorea/n_common.h"
#include "nilorea/n_log.h"
#include "nilorea/n_pth.h"

static int log_level = LOG_ERR;
static int failures = 0;

static const unsigned char checkpoint_pickle[] = {
    0x80, 0x02, 0x7d, 0x71, 0x00, 0x28, 0x58, 0x04, 0x00, 0x00, 0x00, 0x6d,
    0x65, 0x74, 0x61, 0x71, 0x01, 0x7d, 0x71, 0x02, 0x28, 0x58, 0x05, 0x00,
    0x00, 0x00, 0x65, 0x70, 0x6f, 0x63, 0x68, 0x71, 0x03, 0x4b, 0x0c, 0x58,
    0x04, 0x00, 0x00, 0x00, 0x6e, 0x6f, 0x74, 0x65, 0x71, 0x04, 0x58, 0x01,
    0x00, 0x00, 0x00, 0x78, 0x71, 0x05, 0x58, 0x02, 0x00, 0x00, 0x00, 0x6c,
    0x72, 0x71, 0x06, 0x47, 0x3f, 0x50, 0x62, 0x4d, 0xd2, 0xf1, 0xa9, 0xfc,
    0x75, 0x58, 0x0a, 0x00, 0x00, 0x00, 0x73, 0x74, 0x61, 0x74, 0x65, 0x5f,
    0x64, 0x69, 0x63, 0x74, 0x71, 0x07, 0x63, 0x63, 0x6f, 0x6c, 0x6c, 0x65,
    0x63, 0x74, 0x69, 0x6f, 0x6e, 0x73, 0x0a, 0x4f, 0x72, 0x64, 0x65, 0x72,
    0x65, 0x64, 0x44, 0x69, 0x63, 0x74, 0x0a, 0x71, 0x08, 0x29, 0x52, 0x71,
    0x09, 0x28, 0x58, 0x0c, 0x00, 0x00, 0x00, 0x6c, 0x61, 0x79, 0x65, 0x72,
    0x2e, 0x77, 0x65, 0x69, 0x67, 0x68, 0x74, 0x71, 0x0a, 0x63, 0x74, 0x6f,
    0x72, 0x63, 0x68, 0x2e, 0x5f, 0x75, 0x74, 0x69, 0x6c, 0x73, 0x0a, 0x5f,
    0x72, 0x65, 0x62, 0x75, 0x69, 0x6c, 0x64, 0x5f, 0x74, 0x65, 0x6e, 0x73,
    0x6f, 0x72, 0x5f, 0x76, 0x32, 0x0a, 0x71, 0x0b, 0x28, 0x28, 0x58, 0x07,
    0x00, 0x00, 0x00, 0x73, 0x74, 0x6f, 0x72, 0x61, 0x67, 0x65, 0x71, 0x0c,
    0x63, 0x74, 0x6f, 0x72, 0x63, 0x68, 0x0a, 0x46, 0x6c, 0x6f, 0x61, 0x74,
    0x53, 0x74, 0x6f, 0x72, 0x61, 0x67, 0x65, 0x0a, 0x71, 0x0d, 0x58, 0x01,
    0x00, 0x00, 0x00, 0x30, 0x71, 0x0e, 0x58, 0x03, 0x00, 0x00, 0x00, 0x63,
    0x70, 0x75, 0x71, 0x0f, 0x4b, 0x06, 0x74, 0x71, 0x10, 0x51, 0x4b, 0x00,
    0x4b, 0x02, 0x4b, 0x03, 0x86, 0x71, 0x11, 0x4b, 0x03, 0x4b, 0x01, 0x86,
    0x71, 0x12, 0x89, 0x68, 0x08, 0x29, 0x52, 0x71, 0x13, 0x74, 0x71, 0x14,
    0x52, 0x71, 0x15, 0x58, 0x0e, 0x00, 0x00, 0x00, 0x6c, 0x61, 0x79, 0x65,
    0x72, 0x2e, 0x77, 0x65, 0x69, 0x67, 0x68, 0x74, 0x5f, 0x74, 0x71, 0x16,
    0x68, 0x0b, 0x28, 0x68, 0x10, 0x51, 0x4b, 0x00, 0x4b, 0x03, 0x4b, 0x02,
    0x86, 0x71, 0x17, 0x4b, 0x01, 0x4b, 0x03, 0x86, 0x71, 0x18, 0x89, 0x68,
    0x08, 0x29, 0x52, 0x71, 0x19, 0x74, 0x71, 0x1a, 0x52, 0x71, 0x1b, 0x58,
    0x0a, 0x00, 0x00, 0x00, 0x6c, 0x61, 0x79, 0x65, 0x72, 0x2e, 0x62, 0x69,
    0x61, 0x73, 0x71, 0x1c, 0x68, 0x0b, 0x28, 0x68, 0x10, 0x51, 0x4b, 0x04,
    0x4b, 0x02, 0x85, 0x71, 0x1d, 0x4b, 0x01, 0x85, 0x71, 0x1e, 0x89, 0x68,
    0x08, 0x29, 0x52, 0x71, 0x1f, 0x74, 0x71, 0x20, 0x52, 0x71, 0x21, 0x58,
    0x04, 0x00, 0x00, 0x00, 0x68, 0x61, 0x6c, 0x66, 0x71, 0x22, 0x68, 0x0b,
    0x28, 0x28, 0x68, 0x0c, 0x63, 0x74, 0x6f, 0x72, 0x63, 0x68, 0x0a, 0x48,
    0x61, 0x6c, 0x66, 0x53, 0x74, 0x6f, 0x72, 0x61, 0x67, 0x65, 0x0a, 0x71,
    0x23, 0x58, 0x01, 0x00, 0x00, 0x00, 0x31, 0x71, 0x24, 0x68, 0x0f, 0x4b,
    0x04, 0x74, 0x71, 0x25, 0x51, 0x4b, 0x00, 0x4b, 0x04, 0x85, 0x71, 0x26,
    0x68, 0x1e, 0x89, 0x68, 0x08, 0x29, 0x52, 0x71, 0x27, 0x74, 0x71, 0x28,
    0x52, 0x71, 0x29, 0x75, 0x75, 0x2e};

static void expect_true(const char* what, int ok) {
    if (!ok) {
        n_log(LOG_ERR, "FAIL: %s", what);
        failures++;
    } else {
        n_log(LOG_INFO, "ok: %s", what);
    }
}

static int process_args(int argc, char** argv) {
    int opt = 0;
 
    while ((opt = getopt(argc, argv, "V:")) != -1) {
        if (opt == 'V') {
            log_level = atoi(optarg);
        } else {
            return -1;
        }
    }
    return 0;
}

static void put_le(unsigned char* p, unsigned long v, int n) {
    int i = 0;
 
    for (i = 0; i < n; i++) p[i] = (unsigned char)((v >> (8 * i)) & 0xff);
}

typedef struct {
    const char* name;
    const unsigned char* data;
    size_t len;
} MEMBER;

static int write_zip(const char* path, const MEMBER* m, int count) {
    FILE* f = fopen(path, "wb");
    unsigned long offsets[8], crcs[8], at = 0, dir_size = 0;
    unsigned char h[46];
    int i = 0;
 
    if (!f || count > 8) {
        if (f) fclose(f);
        return -1;
    }
    for (i = 0; i < count; i++) {
        size_t n = strlen(m[i].name);
 
        crcs[i] = crc32(crc32(0L, Z_NULL, 0), m[i].data, (uInt)m[i].len);
        memset(h, 0, sizeof(h));
        put_le(h, 0x04034b50ul, 4);
        put_le(h + 4, 20, 2);
        put_le(h + 14, crcs[i], 4);
        put_le(h + 18, (unsigned long)m[i].len, 4);
        put_le(h + 22, (unsigned long)m[i].len, 4);
        put_le(h + 26, (unsigned long)n, 2);
        offsets[i] = at;
        fwrite(h, 1, 30, f);
        fwrite(m[i].name, 1, n, f);
        fwrite(m[i].data, 1, m[i].len, f);
        at += 30 + (unsigned long)n + (unsigned long)m[i].len;
    }
    for (i = 0; i < count; i++) {
        size_t n = strlen(m[i].name);
 
        memset(h, 0, sizeof(h));
        put_le(h, 0x02014b50ul, 4);
        put_le(h + 4, 20, 2);
        put_le(h + 6, 20, 2);
        put_le(h + 16, crcs[i], 4);
        put_le(h + 20, (unsigned long)m[i].len, 4);
        put_le(h + 24, (unsigned long)m[i].len, 4);
        put_le(h + 28, (unsigned long)n, 2);
        put_le(h + 42, offsets[i], 4);
        fwrite(h, 1, 46, f);
        fwrite(m[i].name, 1, n, f);
        dir_size += 46 + (unsigned long)n;
    }
    memset(h, 0, 22);
    put_le(h, 0x06054b50ul, 4);
    put_le(h + 8, (unsigned long)count, 2);
    put_le(h + 10, (unsigned long)count, 2);
    put_le(h + 12, dir_size, 4);
    put_le(h + 16, at, 4);
    fwrite(h, 1, 22, f);
    return fclose(f) == 0 ? 0 : -1;
}

static int values_are(const N_PTH_TENSOR* t, const float* want, size_t n) {
    size_t i = 0;
 
    if (!t || t->count != n) return 0;
    for (i = 0; i < n; i++)
        if (t->data[i] != want[i]) return 0;
    return 1;
}

int main(int argc, char** argv) {
    static const char* path = "ex_pth.tmp.pth";
    static const float f32[6] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f};
    /* 0.5, -1, 2 and 65504, the largest float16, little endian halves. */
    static const unsigned char f16[8] = {0x00, 0x38, 0x00, 0xbc, 0x00, 0x40, 0xff, 0x7b};
    unsigned char f32_bytes[24];
    MEMBER members[3];
    N_PTH* pth = NULL;
    int i = 0;
 
    set_log_level(LOG_ERR);
    if (process_args(argc, argv) != 0) {
        fprintf(stderr, "Usage: %s [-V LOG_LEVEL]\n", argv[0]);
        return 2;
    }
    set_log_level(log_level);
    /* The storage as torch writes it: raw little endian float32. */
    for (i = 0; i < 6; i++) {
        uint32_t bits = 0;
 
        memcpy(&bits, &f32[i], 4);
        put_le(f32_bytes + 4 * i, bits, 4);
    }
    members[0].name = "archive/data.pkl";
    members[0].data = checkpoint_pickle;
    members[0].len = sizeof(checkpoint_pickle);
    members[1].name = "archive/data/0";
    members[1].data = f32_bytes;
    members[1].len = sizeof(f32_bytes);
    members[2].name = "archive/data/1";
    members[2].data = f16;
    members[2].len = sizeof(f16);
    expect_true("a checkpoint is written", write_zip(path, members, 3) == 0);
 
    pth = n_pth_load(path);
    expect_true("it is read", pth != NULL);
    if (pth) {
        const N_PTH_TENSOR* w = n_pth_find(pth, "state_dict.layer.weight");
        const float transposed[6] = {1.0f, 4.0f, 2.0f, 5.0f, 3.0f, 6.0f};
        const float bias[2] = {5.0f, 6.0f};
        const float halves[4] = {0.5f, -1.0f, 2.0f, 65504.0f};
 
        expect_true("four tensors, the metadata left out", pth->count == 4);
        expect_true("a tensor has its path, shape and values",
                    w && w->dims == 2 && w->shape[0] == 2 && w->shape[1] == 3 && values_are(w, f32, 6));
        expect_true("a transposed view is gathered",
                    values_are(n_pth_find(pth, "state_dict.layer.weight_t"), transposed, 6));
        expect_true("an offset view starts where it says",
                    values_are(n_pth_find(pth, "state_dict.layer.bias"), bias, 2));
        expect_true("float16 is converted", values_are(n_pth_find(pth, "state_dict.half"), halves, 4));
        expect_true("a missing name is not found", n_pth_find(pth, "state_dict.nothing") == NULL);
        n_pth_free(&pth);
        expect_true("a freed checkpoint is cleared", pth == NULL);
    }
 
    /* A storage missing: refused, not read past. */
    expect_true("a checkpoint missing a storage is written", write_zip(path, members, 2) == 0);
    pth = n_pth_load(path);
    expect_true("and is refused", pth == NULL);
    n_pth_free(&pth);
    /* A pickle cut short. */
    members[0].len = sizeof(checkpoint_pickle) / 2;
    expect_true("a cut pickle is written", write_zip(path, members, 3) == 0);
    pth = n_pth_load(path);
    expect_true("and is refused", pth == NULL);
    n_pth_free(&pth);
    /* A storage shorter than the tensor that reads it. */
    members[0].len = sizeof(checkpoint_pickle);
    members[1].len = 8;
    expect_true("a short storage is written", write_zip(path, members, 3) == 0);
    pth = n_pth_load(path);
    expect_true("and is refused", pth == NULL);
    n_pth_free(&pth);
    expect_true("a file that is not a checkpoint is refused", n_pth_load("ex_pth.c") == NULL);
    remove(path);
    if (failures) {
        fprintf(stderr, "ex_pth: %d failure(s)\n", failures);
        return 1;
    }
    printf("ex_pth: ok\n");
    return 0;
}