seamless_communication/ggml/examples/unity/model_loader.cpp

#include <string>
#include "model_loader.h"

#define DEBUG_MODEL_LOAD 0

std::ifstream open_ggml_file(const char* fname) {
    printf("%s: loading model from '%s'\n", __func__, fname);

    auto fin = std::ifstream(std::string(fname), std::ios::binary);
    if (!fin) {
        fprintf(stderr, "%s: failed to open '%s'\n", __func__, fname);
        throw std::invalid_argument("failed to open file."); // TODO Merge error message.
    }

    std::uint32_t magic;
    fin.read((char*)&magic, 4);
    if (magic != GGML_FILE_MAGIC) {
        fprintf(stderr, "%s: invalid model file '%s' (bad header %d)\n", __func__, fname, magic);
        throw std::invalid_argument("failed to open file."); // TODO Merge error message.
    }
    return fin;
}

void register_prefix(fairseq2_model &model, const std::string& name) {
    std::size_t i = name.find_last_of('.');
    while(i != std::string::npos && i > 0) {
        std::string prefix = name.substr(0, i);
        auto prev_tensor = model.tensors.find(prefix);
        if (prev_tensor != model.tensors.end()) {
            GGML_ASSERT(prev_tensor->second == nullptr);
        }
        model.tensors[prefix] = nullptr;
        i = name.find_last_of('.', i - 1);
    }
}


std::int64_t
model_loader::load_model_weights(fairseq2_model &model, std::ifstream &fin)
{
    std::int64_t num_tensor = 0;
    std::int64_t ctx_size = 0;
    fin.read((char*) &num_tensor, sizeof(num_tensor));
    fin.read((char*) &ctx_size, sizeof(ctx_size));

    struct ggml_init_params params = {
        /*.mem_size   =*/ static_cast<std::size_t>(ctx_size),
        /*.mem_buffer =*/ NULL,
        /*.no_alloc   =*/ false,
    };
    model.tensors_ctx = ggml_init(params);

    size_t model_size = 0;
    for (int i = 0; i < num_tensor; ++i) {
        std::string name = get_name(fin);
        if (name.length() == 0)
            break;
        auto tensor = load_tensor_value(fin, model.tensors_ctx);
        if (tensor == nullptr) {
            // Abort in case of error, the input stream is corrupted at this point.
            printf("Error while reading tensor %s\n", name.c_str() );
            throw std::invalid_argument("Error while reading tensor from file.");
        }
        register_prefix(model, name);
        ggml_set_name(tensor, name.c_str());
        model.tensors[name] = tensor;
        if (DEBUG_MODEL_LOAD) {
            printf("%s [%5ld, %5ld], type = %6s, %6.2f MB, %9zu bytes\n", name.c_str(), tensor->ne[0], tensor->ne[1], ggml_type_name(tensor->type), ggml_nbytes(tensor)/1024.0/1024.0, ggml_nbytes(tensor));
        }
        model_size += ggml_nbytes(tensor);
    }

    double mb = 1024.0 * 1024.0;
    printf("%s: model size  = %8.2f MB, memory used = %8.2f MB, memory reserved = %8.2f \n",
        __func__,
        model_size / mb,
        ggml_used_mem(model.tensors_ctx) / mb,
        ctx_size / mb
    );

    return ctx_size;
}

void assert_endianness() {
    union {
        unsigned int i;
        char c[4];
    } un;
    un.i = 0x12345678;

    if (un.c[0] == 0x78 && un.c[3] == 0x12) {
        printf("little-endian\n");
    }
    else if (un.c[0] == 0x12 && un.c[3] == 0x78) {
        printf("big-endian\n");
        GGML_ASSERT(false); // model_loader.cpp assumes the system is little-endian
    }
    else {
        printf("unknown-endian\n");
        GGML_ASSERT(false); // model_loader.cpp assumes the system is little-endian
    }
}


void model_loader::load_hparams(std::unordered_map<std::string, std::int64_t>& hparams, std::ifstream &fin)
{
    std::int64_t num_params = 0;
    fin.read(reinterpret_cast<char*>(&num_params), sizeof num_params);
    GGML_ASSERT(fin.gcount() == 8);

    hparams.reserve(num_params);

    std::int64_t value;
    for (int i = 0; i < num_params; ++i) {
        std::string name = get_name(fin);
        if (name.length() == 0)
            break;
        fin.read((char*) &value, sizeof(value));
        hparams[name] = value;
    }
}

ggml_tensor* load_tensor_value(std::ifstream &fin, ggml_context* ctx)
{
    int32_t n_dims = 0;
    int32_t raw_type = 0;

    fin.read(reinterpret_cast<char *>(&n_dims), sizeof(n_dims));
    fin.read(reinterpret_cast<char *>(&raw_type),  sizeof(raw_type));
    ggml_type type = ggml_type(raw_type);

    if (n_dims <= 0 || n_dims > GGML_MAX_DIMS || raw_type < 0 || raw_type > GGML_TYPE_COUNT) {
        return nullptr;
    }
    int64_t ne[4] = {1, 1, 1, 1};
    for (int i = 0; i < n_dims; ++i) {
        fin.read(reinterpret_cast<char *>(&ne[i]), sizeof(ne[i]));
    }

    ggml_tensor* tensor = ggml_new_tensor(ctx, type, n_dims, ne);
    fin.read(reinterpret_cast<char *>(tensor->data), ggml_nbytes(tensor));
    return tensor;
}

std::string
model_loader::get_name(std::ifstream& fin)
{
    std::uint32_t length = 0;
    fin.read(reinterpret_cast<char *>(&length), sizeof(length));
    if (length == 0)
        return "";

    std::string name(length, 0);
    fin.read(&name[0], length);

    return name;
}

extern "C" int load_fairseq2_ggml_file(fairseq2_model& model, const char* fname) {
    model_loader loader;
    assert_endianness();
    auto fin = open_ggml_file(fname);
    loader.load_hparams(model.hparams, fin);
    loader.load_hparams(model.layer_config, fin);
    loader.load_model_weights(model, fin);
    return 0;
}