char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Pepper Grayllama : add missing headers (#28566)ca86fb222

char/aliceai-t5-moe
56.7 KiB1707 linesraw
1#include "ggml.h"
2#include "ggml-backend.h"
3#include "ggml-impl.h"
4#include "gguf.h"
5
6#include <cinttypes>
7#include <cstddef>
8#include <cstdint>
9#include <cstdio>
10#include <cstdlib>
11#include <cstring>
12#include <cerrno>
13#include <map>
14#include <new>
15#include <stdexcept>
16#include <string>
17#include <vector>
18
19#define GGUF_MAX_STRING_LENGTH  (1024*1024*1024)
20#define GGUF_MAX_ARRAY_ELEMENTS (1024*1024*1024)
21
22#ifdef _WIN32
23#    define gguf_ftell _ftelli64
24#    define gguf_fseek _fseeki64
25#else
26#    define gguf_ftell ftello
27#    define gguf_fseek fseeko
28#endif
29
30template <typename T>
31struct type_to_gguf_type;
32
33template <>
34struct type_to_gguf_type<uint8_t> {
35    static constexpr enum gguf_type value = GGUF_TYPE_UINT8;
36};
37
38template <>
39struct type_to_gguf_type<int8_t> {
40    static constexpr enum gguf_type value = GGUF_TYPE_INT8;
41};
42
43template <>
44struct type_to_gguf_type<uint16_t> {
45    static constexpr enum gguf_type value = GGUF_TYPE_UINT16;
46};
47
48template <>
49struct type_to_gguf_type<int16_t> {
50    static constexpr enum gguf_type value = GGUF_TYPE_INT16;
51};
52
53template <>
54struct type_to_gguf_type<uint32_t> {
55    static constexpr enum gguf_type value = GGUF_TYPE_UINT32;
56};
57
58template <>
59struct type_to_gguf_type<int32_t> {
60    static constexpr enum gguf_type value = GGUF_TYPE_INT32;
61};
62
63template <>
64struct type_to_gguf_type<float> {
65    static constexpr enum gguf_type value = GGUF_TYPE_FLOAT32;
66};
67
68template <>
69struct type_to_gguf_type<bool> {
70    static constexpr enum gguf_type value = GGUF_TYPE_BOOL;
71};
72
73template <>
74struct type_to_gguf_type<std::string> {
75    static constexpr enum gguf_type value = GGUF_TYPE_STRING;
76};
77
78template <>
79struct type_to_gguf_type<uint64_t> {
80    static constexpr enum gguf_type value = GGUF_TYPE_UINT64;
81};
82
83template <>
84struct type_to_gguf_type<int64_t> {
85    static constexpr enum gguf_type value = GGUF_TYPE_INT64;
86};
87
88template <>
89struct type_to_gguf_type<double> {
90    static constexpr enum gguf_type value = GGUF_TYPE_FLOAT64;
91};
92
93static const std::map<gguf_type, size_t> GGUF_TYPE_SIZE = {
94    {GGUF_TYPE_UINT8,   sizeof(uint8_t)},
95    {GGUF_TYPE_INT8,    sizeof(int8_t)},
96    {GGUF_TYPE_UINT16,  sizeof(uint16_t)},
97    {GGUF_TYPE_INT16,   sizeof(int16_t)},
98    {GGUF_TYPE_UINT32,  sizeof(uint32_t)},
99    {GGUF_TYPE_INT32,   sizeof(int32_t)},
100    {GGUF_TYPE_FLOAT32, sizeof(float)},
101    {GGUF_TYPE_BOOL,    sizeof(int8_t)},
102    {GGUF_TYPE_STRING,  0}, // undefined
103    {GGUF_TYPE_ARRAY,   0}, // undefined
104    {GGUF_TYPE_UINT64,  sizeof(uint64_t)},
105    {GGUF_TYPE_INT64,   sizeof(int64_t)},
106    {GGUF_TYPE_FLOAT64, sizeof(double)},
107};
108static_assert(GGUF_TYPE_COUNT == 13, "GGUF_TYPE_COUNT != 13");
109
110static const std::map<gguf_type, const char *> GGUF_TYPE_NAME = {
111    {GGUF_TYPE_UINT8,   "u8"},
112    {GGUF_TYPE_INT8,    "i8"},
113    {GGUF_TYPE_UINT16,  "u16"},
114    {GGUF_TYPE_INT16,   "i16"},
115    {GGUF_TYPE_UINT32,  "u32"},
116    {GGUF_TYPE_INT32,   "i32"},
117    {GGUF_TYPE_FLOAT32, "f32"},
118    {GGUF_TYPE_BOOL,    "bool"},
119    {GGUF_TYPE_STRING,  "str"},
120    {GGUF_TYPE_ARRAY,   "arr"},
121    {GGUF_TYPE_UINT64,  "u64"},
122    {GGUF_TYPE_INT64,   "i64"},
123    {GGUF_TYPE_FLOAT64, "f64"},
124};
125static_assert(GGUF_TYPE_COUNT == 13, "GGUF_TYPE_COUNT != 13");
126
127size_t gguf_type_size(enum gguf_type type) {
128    auto it = GGUF_TYPE_SIZE.find(type);
129    return it == GGUF_TYPE_SIZE.end() ? 0 : it->second;
130}
131
132struct gguf_kv {
133    std::string key;
134
135    bool is_array;
136    enum gguf_type type;
137
138    std::vector<int8_t>      data;
139    std::vector<std::string> data_string;
140
141    template <typename T>
142    gguf_kv(const std::string & key, const T value)
143            : key(key), is_array(false), type(type_to_gguf_type<T>::value) {
144        GGML_ASSERT(!key.empty());
145        data.resize(sizeof(T));
146        memcpy(data.data(), &value, sizeof(T));
147    }
148
149    template <typename T>
150    gguf_kv(const std::string & key, const std::vector<T> & value)
151            : key(key), is_array(true), type(type_to_gguf_type<T>::value) {
152        GGML_ASSERT(!key.empty());
153        data.resize(value.size()*sizeof(T));
154        for (size_t i = 0; i < value.size(); ++i) {
155            const T tmp = value[i];
156            memcpy(data.data() + i*sizeof(T), &tmp, sizeof(T));
157        }
158    }
159
160    gguf_kv(const std::string & key, const std::string & value)
161            : key(key), is_array(false), type(GGUF_TYPE_STRING) {
162        GGML_ASSERT(!key.empty());
163        data_string.push_back(value);
164    }
165
166    gguf_kv(const std::string & key, const std::vector<std::string> & value)
167            : key(key), is_array(true), type(GGUF_TYPE_STRING) {
168        GGML_ASSERT(!key.empty());
169        data_string = value;
170    }
171
172    const std::string & get_key() const {
173        return key;
174    }
175
176    const enum gguf_type & get_type() const {
177        return type;
178    }
179
180    size_t get_ne() const {
181        if (type == GGUF_TYPE_STRING) {
182            const size_t ne = data_string.size();
183            GGML_ASSERT(is_array || ne == 1);
184            return ne;
185        }
186        const size_t type_size = gguf_type_size(type);
187        GGML_ASSERT(data.size() % type_size == 0);
188        const size_t ne = data.size() / type_size;
189        GGML_ASSERT(is_array || ne == 1);
190        return ne;
191    }
192
193    template <typename T>
194    const T & get_val(const size_t i = 0) const {
195        GGML_ASSERT(type_to_gguf_type<T>::value == type);
196        if constexpr (std::is_same<T, std::string>::value) {
197            GGML_ASSERT(data_string.size() >= i+1);
198            return data_string[i];
199        }
200        const size_t type_size = gguf_type_size(type);
201        GGML_ASSERT(data.size() % type_size == 0);
202        GGML_ASSERT(data.size() >= (i+1)*type_size);
203        return reinterpret_cast<const T *>(data.data())[i];
204    }
205
206    void cast(const enum gguf_type new_type) {
207        const size_t new_type_size = gguf_type_size(new_type);
208        GGML_ASSERT(data.size() % new_type_size == 0);
209        type = new_type;
210    }
211};
212
213struct gguf_tensor_info {
214    struct ggml_tensor t; // for holding the equivalent info
215    uint64_t offset;      // offset from start of `data`, must be a multiple of `ALIGNMENT`
216};
217
218struct gguf_context {
219    uint32_t version = GGUF_VERSION;
220
221    std::vector<struct gguf_kv> kv;
222    std::vector<struct gguf_tensor_info> info;
223
224    size_t alignment = GGUF_DEFAULT_ALIGNMENT;
225    size_t offset    = 0; // offset of `data` from beginning of file
226    size_t size      = 0; // size of `data` in bytes
227
228    void * data = nullptr;
229};
230
231struct gguf_reader {
232    gguf_reader(
233            gguf_reader_callback_t callback,
234            void * userdata,
235            size_t max_chunk_read,
236            uint64_t data_offset = 0,
237            uint64_t nbytes_remain = 0)
238        : callback(callback),
239          userdata(userdata),
240          max_chunk_read(max_chunk_read),
241          data_offset(data_offset),
242          nbytes_remain(nbytes_remain) {
243        GGML_ASSERT(max_chunk_read > 0);
244    }
245
246    // helper for remaining bytes in a file
247    static uint64_t file_remain(FILE * file) {
248        const int64_t cur = gguf_ftell(file);
249        if (cur < 0) {
250            return 0;
251        }
252        if (gguf_fseek(file, 0, SEEK_END) != 0) {
253            gguf_fseek(file, cur, SEEK_SET);
254
255            return 0;
256        }
257        const int64_t end = gguf_ftell(file);
258        if (end < 0) {
259            gguf_fseek(file, cur, SEEK_SET);
260
261            return 0;
262        }
263        gguf_fseek(file, cur, SEEK_SET);
264        return static_cast<uint64_t>(end - cur);
265    }
266
267    template <typename T>
268    bool read(T & dst) const {
269        const size_t size = sizeof(dst);
270        if (size > nbytes_remain) {
271            return false;
272        }
273        return read_raw(&dst, size) == size;
274    }
275
276    template <typename T>
277    bool read(std::vector<T> & dst, const size_t n) const {
278        if (n > GGUF_MAX_ARRAY_ELEMENTS) {
279            return false;
280        }
281        if constexpr (std::is_same<T, std::string>::value) {
282            // strings are prefixed with their length, so we need to account for that
283            if (n > SIZE_MAX / sizeof(uint64_t)) {
284                return false;
285            }
286            if (nbytes_remain < n * sizeof(uint64_t)) {
287                return false;
288            }
289        } else {
290            if (n > SIZE_MAX / sizeof(T)) {
291                return false;
292            }
293            if (nbytes_remain < n * sizeof(T)) {
294                return false;
295            }
296        }
297        dst.resize(n);
298        for (size_t i = 0; i < dst.size(); ++i) {
299            if constexpr (std::is_same<T, bool>::value) {
300                bool tmp;
301                if (!read(tmp)) {
302                    return false;
303                }
304                dst[i] = tmp;
305            } else {
306                if (!read(dst[i])) {
307                    return false;
308                }
309            }
310        }
311        return true;
312    }
313
314    bool read(bool & dst) const {
315        int8_t tmp = -1;
316        if (!read(tmp)) {
317            return false;
318        }
319        dst = tmp != 0;
320        return true;
321    }
322
323    bool read(enum ggml_type & dst) const {
324        int32_t tmp = -1;
325        if (!read(tmp)) {
326            return false;
327        }
328        dst = ggml_type(tmp);
329        return true;
330    }
331
332    bool read(enum gguf_type & dst) const {
333        int32_t tmp = -1;
334        if (!read(tmp)) {
335            return false;
336        }
337        dst = gguf_type(tmp);
338        return true;
339    }
340
341    bool read(std::string & dst) const {
342        uint64_t size = 0;
343        if (!read(size)) {
344            return false;
345        }
346        if (size > GGUF_MAX_STRING_LENGTH) {
347            GGML_LOG_ERROR("%s: string length %" PRIu64 " exceeds maximum %" PRIu64 "\n", __func__, size, (uint64_t) GGUF_MAX_STRING_LENGTH);
348            return false;
349        }
350        if (size > nbytes_remain) {
351            GGML_LOG_ERROR("%s: string length %" PRIu64 " exceeds remaining file size %" PRIu64 " bytes\n", __func__, size, nbytes_remain);
352            return false;
353        }
354        dst.resize(static_cast<size_t>(size));
355        return read_raw(dst.data(), static_cast<size_t>(size)) == size;
356    }
357
358    bool read(void * dst, const size_t size) const {
359        if (size > nbytes_remain) {
360            return false;
361        }
362        return read_raw(dst, size) == size;
363    }
364
365    uint64_t tell() const {
366        return data_offset;
367    }
368
369    bool seek(uint64_t absolute_offset) const {
370        const uint64_t end_offset = uint64_t(data_offset) + nbytes_remain;
371        if (absolute_offset > end_offset) {
372            return false;
373        }
374
375        data_offset = absolute_offset;
376        nbytes_remain = end_offset - absolute_offset;
377
378        return true;
379    }
380
381private:
382    size_t read_raw(void * dst, size_t size) const {
383        if (callback == nullptr || size == 0) {
384            return 0;
385        }
386
387        uint8_t * data = static_cast<uint8_t *>(dst);
388        size_t total_nread = 0;
389        bool reached_eof = false;
390
391        while (total_nread < size) {
392            const size_t chunk_size = std::min(max_chunk_read, size - total_nread);
393            if (data_offset + total_nread < data_offset) {
394                break;
395            }
396            const size_t nread = callback(userdata, static_cast<void *>(data + total_nread), data_offset + total_nread, chunk_size);
397            total_nread += nread;
398            if (nread != chunk_size) {
399                reached_eof = true;
400                break;
401            }
402        }
403
404        data_offset += total_nread;
405        GGML_ASSERT(total_nread <= nbytes_remain);
406        nbytes_remain -= total_nread;
407
408        if (reached_eof) {
409            nbytes_remain = 0;
410        }
411
412        return total_nread;
413    }
414
415    gguf_reader_callback_t callback = nullptr;
416    void * userdata = nullptr;
417    size_t max_chunk_read = 0;
418    mutable uint64_t data_offset = 0;
419    mutable uint64_t nbytes_remain = 0;
420};
421
422struct gguf_context * gguf_init_empty(void) {
423    return new gguf_context;
424}
425
426template<typename T>
427bool gguf_read_emplace_helper(const struct gguf_reader & gr, std::vector<struct gguf_kv> & kv, const std::string & key, const bool is_array, const size_t n) {
428    if (is_array) {
429        std::vector<T> value;
430        try {
431            if (!gr.read(value, n)) {
432                return false;
433            }
434        } catch (std::length_error &) {
435            GGML_LOG_ERROR("%s: encountered length_error while reading value for key '%s'\n", __func__, key.c_str());
436            return false;
437        } catch (std::bad_alloc &) {
438            GGML_LOG_ERROR("%s: encountered bad_alloc error while reading value for key '%s'\n", __func__, key.c_str());
439            return false;
440        }
441        kv.emplace_back(key, value);
442    } else {
443        T value;
444        if (!gr.read(value)) {
445            return false;
446        }
447        kv.emplace_back(key, value);
448    }
449    return true;
450}
451
452static struct gguf_context * gguf_init_from_reader(const struct gguf_reader & gr, struct gguf_init_params params) {
453    struct gguf_context * ctx = new gguf_context;
454
455    bool ok = true;
456
457    // file magic
458    {
459        std::vector<char> magic;
460        ok = ok && gr.read(magic, 4);
461
462        if (!ok) {
463            GGML_LOG_ERROR("%s: failed to read magic\n", __func__);
464            gguf_free(ctx);
465            return nullptr;
466        }
467
468        for (uint32_t i = 0; i < magic.size(); i++) {
469            if (magic[i] != GGUF_MAGIC[i]) {
470                char c0 = isprint(magic[0]) ? magic[0] : '?';
471                char c1 = isprint(magic[1]) ? magic[1] : '?';
472                char c2 = isprint(magic[2]) ? magic[2] : '?';
473                char c3 = isprint(magic[3]) ? magic[3] : '?';
474                GGML_LOG_ERROR("%s: invalid magic characters: '%c%c%c%c', expected 'GGUF'\n", __func__, c0, c1, c2, c3);
475                gguf_free(ctx);
476                return nullptr;
477            }
478        }
479    }
480
481    // header
482    int64_t n_kv      = 0;
483    int64_t n_tensors = 0;
484
485    if (ok && gr.read(ctx->version)) {
486        if (ok && ctx->version == 0) {
487            GGML_LOG_ERROR("%s: bad GGUF version: %" PRIu32 "\n", __func__, ctx->version);
488            ok = false;
489        }
490
491        /*
492         * bit layout is different when reading non-native endian models.
493         * assuming that the GGUF version is 3, the non-native endian model
494         * would read it as 0x30000000. we can use the AND operation against
495         * the last 4 hexadecimal digits to check if the model is the same
496         * endianness as the host system.
497        */
498        if (ok && (ctx->version & 0x0000FFFF) == 0x00000000) {
499            GGML_LOG_ERROR("%s: failed to load model: this GGUF file version %" PRIu32 " is extremely large, is there a mismatch between the host and model endianness?\n", __func__, ctx->version);
500            ok = false;
501        }
502
503        if (ok && ctx->version == 1) {
504            GGML_LOG_ERROR("%s: GGUFv1 is no longer supported, please use a more up-to-date version\n", __func__);
505            ok = false;
506        }
507        if (ok && ctx->version > GGUF_VERSION) {
508            GGML_LOG_ERROR("%s: this GGUF file is version %" PRIu32 " but this software only supports up to version %d\n",
509                __func__, ctx->version, GGUF_VERSION);
510            ok = false;
511        }
512    } else {
513        ok = false;
514    }
515
516    if (ok && gr.read(n_tensors)) {
517        static_assert(sizeof(size_t) <= 8 && sizeof(gguf_tensor_info) >= 2, "int64_t insufficient for indexing");
518        if (n_tensors < 0 || n_tensors > int64_t(SIZE_MAX/sizeof(gguf_tensor_info))) {
519            GGML_LOG_ERROR("%s: number of tensors is %" PRIi64 " but must be in [0, %zu]\n",
520                __func__, n_tensors, SIZE_MAX/sizeof(gguf_tensor_info));
521            ok = false;
522        }
523    } else {
524        ok = false;
525    }
526
527    if (ok && gr.read(n_kv)) {
528        static_assert(sizeof(size_t) <= 8 && sizeof(gguf_tensor_info) >= 2, "int64_t insufficient for indexing");
529        if (n_kv < 0 || n_kv > int64_t(SIZE_MAX/sizeof(gguf_kv))) {
530            GGML_LOG_ERROR("%s: number of key value pairs is %" PRIi64 " but must be in [0, %zu]\n",
531                    __func__, n_kv, SIZE_MAX/sizeof(gguf_kv));
532            ok = false;
533        }
534    } else {
535        ok = false;
536    }
537
538    if (!ok) {
539        GGML_LOG_ERROR("%s: failed to read header\n", __func__);
540        gguf_free(ctx);
541        return nullptr;
542    }
543
544    // KV pairs
545    {
546        for (int64_t i = 0; ok && i < n_kv; ++i) {
547            std::string key;
548            gguf_type   type     = gguf_type(-1);
549            bool        is_array = false;
550            uint64_t    n        = 1;
551
552            try {
553                ok = ok && gr.read(key);
554            } catch (std::length_error &) {
555                GGML_LOG_ERROR("%s: encountered length_error while reading key %" PRIi64 "\n", __func__, i);
556                ok = false;
557            } catch (std::bad_alloc &) {
558                GGML_LOG_ERROR("%s: encountered bad_alloc error while reading key %" PRIi64 "\n", __func__, i);
559                ok = false;
560            }
561            if (ok && key.empty()) {
562                GGML_LOG_ERROR("%s: key %" PRIi64 " is empty\n", __func__, i);
563                ok = false;
564            }
565            for (size_t j = 0; ok && j < ctx->kv.size(); ++j) {
566                if (key == ctx->kv[j].key) {
567                    GGML_LOG_ERROR("%s: duplicate key '%s' for tensors %zu and %" PRIi64 " \n", __func__, key.c_str(), j, i);
568                    ok = false;
569                }
570            }
571            if (!ok) {
572                break;
573            }
574
575            ok = ok && gr.read(type);
576            if (type == GGUF_TYPE_ARRAY) {
577                is_array = true;
578                ok = ok && gr.read(type);
579                ok = ok && gr.read(n);
580            }
581            if (!ok) {
582                break;
583            }
584
585            switch (type) {
586                case GGUF_TYPE_UINT8:   ok = ok && gguf_read_emplace_helper<uint8_t>    (gr, ctx->kv, key, is_array, n); break;
587                case GGUF_TYPE_INT8:    ok = ok && gguf_read_emplace_helper<int8_t>     (gr, ctx->kv, key, is_array, n); break;
588                case GGUF_TYPE_UINT16:  ok = ok && gguf_read_emplace_helper<uint16_t>   (gr, ctx->kv, key, is_array, n); break;
589                case GGUF_TYPE_INT16:   ok = ok && gguf_read_emplace_helper<int16_t>    (gr, ctx->kv, key, is_array, n); break;
590                case GGUF_TYPE_UINT32:  ok = ok && gguf_read_emplace_helper<uint32_t>   (gr, ctx->kv, key, is_array, n); break;
591                case GGUF_TYPE_INT32:   ok = ok && gguf_read_emplace_helper<int32_t>    (gr, ctx->kv, key, is_array, n); break;
592                case GGUF_TYPE_FLOAT32: ok = ok && gguf_read_emplace_helper<float>      (gr, ctx->kv, key, is_array, n); break;
593                case GGUF_TYPE_BOOL:    ok = ok && gguf_read_emplace_helper<bool>       (gr, ctx->kv, key, is_array, n); break;
594                case GGUF_TYPE_STRING:  ok = ok && gguf_read_emplace_helper<std::string>(gr, ctx->kv, key, is_array, n); break;
595                case GGUF_TYPE_UINT64:  ok = ok && gguf_read_emplace_helper<uint64_t>   (gr, ctx->kv, key, is_array, n); break;
596                case GGUF_TYPE_INT64:   ok = ok && gguf_read_emplace_helper<int64_t>    (gr, ctx->kv, key, is_array, n); break;
597                case GGUF_TYPE_FLOAT64: ok = ok && gguf_read_emplace_helper<double>     (gr, ctx->kv, key, is_array, n); break;
598                case GGUF_TYPE_ARRAY:
599                default:
600                    {
601                        GGML_LOG_ERROR("%s: key '%s' has invalid GGUF type %d\n", __func__, key.c_str(), type);
602                        ok = false;
603                    } break;
604            }
605        }
606
607        if (!ok) {
608            GGML_LOG_ERROR("%s: failed to read key-value pairs\n", __func__);
609            gguf_free(ctx);
610            return nullptr;
611        }
612        GGML_ASSERT(int64_t(ctx->kv.size()) == n_kv);
613
614        const int alignment_idx = gguf_find_key(ctx, GGUF_KEY_GENERAL_ALIGNMENT);
615        if (alignment_idx != -1 && gguf_get_kv_type(ctx, alignment_idx) != GGUF_TYPE_UINT32) {
616            GGML_LOG_ERROR("%s: key '%s' must be of type %s but is %s\n",
617                __func__, GGUF_KEY_GENERAL_ALIGNMENT, gguf_type_name(GGUF_TYPE_UINT32),
618                gguf_type_name(gguf_get_kv_type(ctx, alignment_idx)));
619            gguf_free(ctx);
620            return nullptr;
621        }
622        ctx->alignment = alignment_idx == -1 ? GGUF_DEFAULT_ALIGNMENT : gguf_get_val_u32(ctx, alignment_idx);
623
624        if (ctx->alignment == 0 || (ctx->alignment & (ctx->alignment - 1)) != 0) {
625            GGML_LOG_ERROR("%s: alignment %zu is not a power of 2\n", __func__, ctx->alignment);
626            gguf_free(ctx);
627            return nullptr;
628        }
629    }
630
631    // read the tensor info
632    for (int64_t i = 0; ok && i < n_tensors; ++i) {
633        struct gguf_tensor_info info;
634
635        // tensor name
636        {
637            std::string name;
638            try {
639                ok = ok && gr.read(name);
640            } catch (std::length_error &) {
641                GGML_LOG_ERROR("%s: encountered length_error while reading tensor name %" PRIi64 "\n", __func__, i);
642                ok = false;
643            } catch (std::bad_alloc &) {
644                GGML_LOG_ERROR("%s: encountered bad_alloc error while reading tensor name %" PRIi64 "\n", __func__, i);
645                ok = false;
646            }
647            if (name.length() >= GGML_MAX_NAME) {
648                GGML_LOG_ERROR("%s: tensor name %" PRIi64 " is too long: %zu >= %d\n", __func__, i, name.length(), GGML_MAX_NAME);
649                ok = false;
650                break;
651            }
652            ggml_set_name(&info.t, name.c_str());
653
654            // make sure there are no duplicate tensor names
655            for (int64_t j = 0; ok && j < i; ++j) {
656                if (strcmp(info.t.name, ctx->info[j].t.name) == 0) {
657                    GGML_LOG_ERROR("%s: duplicate tensor name '%s' for tensors %" PRIi64 " and %" PRIi64 "\n", __func__, info.t.name, j, i);
658                    ok = false;
659                    break;
660                }
661            }
662        }
663        if (!ok) {
664            break;
665        }
666
667        // tensor shape
668        {
669            uint32_t n_dims = 0;
670            ok = ok && gr.read(n_dims);
671            if (n_dims > GGML_MAX_DIMS) {
672                GGML_LOG_ERROR("%s: tensor '%s' has invalid number of dimensions: %" PRIu32 " > %" PRIu32 "\n",
673                    __func__, info.t.name, n_dims, GGML_MAX_DIMS);
674                ok = false;
675                break;
676            }
677            for (uint32_t j = 0; ok && j < GGML_MAX_DIMS; ++j) {
678                info.t.ne[j] = 1;
679                if (j < n_dims) {
680                    ok = ok && gr.read(info.t.ne[j]);
681                }
682
683                // check that all ne are non-negative
684                if (info.t.ne[j] < 0) {
685                    GGML_LOG_ERROR("%s: tensor '%s' dimension %" PRIu32 " has invalid number of elements: %" PRIi64 " < 0\n",
686                        __func__, info.t.name, j, info.t.ne[j]);
687                    ok = false;
688                    break;
689                }
690            }
691
692            // check that the total number of elements is representable
693            // (a zero-element tensor is trivially representable; the guard also avoids a division by zero below)
694            if (ok && ggml_nelements(&info.t) > 0 &&
695                ((INT64_MAX/info.t.ne[1] <= info.t.ne[0]) ||
696                 (INT64_MAX/info.t.ne[2] <= info.t.ne[0]*info.t.ne[1]) ||
697                 (INT64_MAX/info.t.ne[3] <= info.t.ne[0]*info.t.ne[1]*info.t.ne[2]))) {
698
699                GGML_LOG_ERROR("%s: total number of elements in tensor '%s' with shape "
700                    "(%" PRIi64 ", %" PRIi64 ", %" PRIi64 ", %" PRIi64 ") is >= %" PRIi64 "\n",
701                    __func__, info.t.name, info.t.ne[0], info.t.ne[1], info.t.ne[2], info.t.ne[3], INT64_MAX);
702                ok = false;
703                break;
704            }
705        }
706        if (!ok) {
707            break;
708        }
709
710        // tensor type
711        {
712            ok = ok && gr.read(info.t.type);
713
714            // check that tensor type is within defined range
715            if (info.t.type < 0 || info.t.type >= GGML_TYPE_COUNT) {
716                GGML_LOG_ERROR("%s: tensor '%s' has invalid ggml type %d. should be in [0, %d)\n",
717                    __func__, info.t.name, info.t.type, GGML_TYPE_COUNT);
718                ok = false;
719                break;
720            }
721            const size_t  type_size = ggml_type_size(info.t.type);
722            const int64_t blck_size = ggml_blck_size(info.t.type);
723
724            // check that row size is divisible by block size
725            if (blck_size == 0 || info.t.ne[0] % blck_size != 0) {
726                GGML_LOG_ERROR("%s: tensor '%s' of type %d (%s) has %" PRId64 " elements per row, "
727                    "not a multiple of block size (%" PRId64 ")\n",
728                    __func__, info.t.name, (int) info.t.type, ggml_type_name(info.t.type), info.t.ne[0], blck_size);
729                ok = false;
730                break;
731            }
732
733            // check that the size of the tensor in bytes is representable
734            if (ok && uint64_t(ggml_nelements(&info.t)/ggml_blck_size(info.t.type)) > SIZE_MAX/ggml_type_size(info.t.type)) {
735                GGML_LOG_ERROR("%s: tensor '%s' with shape (%" PRIi64 ", %" PRIi64 ", %" PRIi64 ", %" PRIi64 ") has a size in bytes > %zu\n",
736                    __func__, info.t.name, info.t.ne[0], info.t.ne[1], info.t.ne[2], info.t.ne[3], SIZE_MAX);
737                ok = false;
738                break;
739            }
740
741            // calculate byte offsets given the tensor shape and type
742            info.t.nb[0] = type_size;
743            info.t.nb[1] = info.t.nb[0]*(info.t.ne[0]/blck_size);
744            for (int j = 2; j < GGML_MAX_DIMS; ++j) {
745                info.t.nb[j] = info.t.nb[j - 1]*info.t.ne[j - 1];
746            }
747        }
748        if (!ok) {
749            break;
750        }
751
752        // tensor data offset within buffer
753        ok = ok && gr.read(info.offset);
754
755        ctx->info.push_back(info);
756    }
757
758    if (!ok) {
759        GGML_LOG_ERROR("%s: failed to read tensor info\n", __func__);
760        gguf_free(ctx);
761        return nullptr;
762    }
763    GGML_ASSERT(int64_t(ctx->info.size()) == n_tensors);
764
765    // we require the data section to be aligned, so take into account any padding
766    if (n_tensors > 0 && !gr.seek(GGML_PAD(gr.tell(), ctx->alignment))) {
767        GGML_LOG_ERROR("%s: failed to seek to beginning of data section\n", __func__);
768        gguf_free(ctx);
769        return nullptr;
770    }
771
772    // store the current file offset - this is where the data section starts
773    ctx->offset = gr.tell();
774
775    // compute the total size of the data section, taking into account the alignment
776    {
777        ctx->size = 0;
778        for (size_t i = 0; i < ctx->info.size(); ++i) {
779            const gguf_tensor_info & ti = ctx->info[i];
780            if (ti.offset != ctx->size) {
781                GGML_LOG_ERROR("%s: tensor '%s' has offset %" PRIu64 ", expected %zu\n",
782                    __func__, ti.t.name, ti.offset, ctx->size);
783                GGML_LOG_ERROR("%s: failed to read tensor data\n", __func__);
784                gguf_free(ctx);
785                return nullptr;
786            }
787            size_t padded_size = GGML_PAD(ggml_nbytes(&ti.t), ctx->alignment);
788            if (SIZE_MAX - ctx->size < padded_size) {
789                GGML_LOG_ERROR("%s: tensor '%s' size overflow, cannot accumulate size %zu + %zu\n",
790                    __func__, ti.t.name, ctx->size, padded_size);
791                gguf_free(ctx);
792                return nullptr;
793            }
794            ctx->size += padded_size;
795        }
796    }
797
798    // load the tensor data only if requested
799    if (params.ctx != nullptr) {
800        // if the provided gguf_context is no_alloc, then we create "empty" tensors and do not read the binary blob
801        // otherwise, we load the binary blob into the created ggml_context as well, and point the "data" members of
802        //   the ggml_tensor structs to the appropriate locations in the binary blob
803
804        // compute the exact size needed for the new ggml_context
805        size_t mem_size = 0;
806        if (params.no_alloc) {
807            if (n_tensors != 0 && SIZE_MAX / n_tensors < ggml_tensor_overhead()) {
808                GGML_LOG_ERROR("%s: memory size overflow while allocating ggml context\n", __func__);
809                gguf_free(ctx);
810                return nullptr;
811            }
812
813            const size_t overhead = n_tensors * ggml_tensor_overhead();
814
815            mem_size = overhead;
816        } else {
817            if ((n_tensors + 1) != 0 && SIZE_MAX / (n_tensors + 1) < ggml_tensor_overhead()) {
818                GGML_LOG_ERROR("%s: memory size overflow while allocating ggml context\n", __func__);
819                gguf_free(ctx);
820                return nullptr;
821            }
822
823            const size_t overhead = (n_tensors + 1) * ggml_tensor_overhead();
824
825            if (SIZE_MAX - overhead < ctx->size) {
826                GGML_LOG_ERROR("%s: memory size overflow while allocating ggml context\n", __func__);
827                gguf_free(ctx);
828                return nullptr;
829            }
830
831            mem_size = overhead + ctx->size;
832        }
833
834        struct ggml_init_params pdata = {
835            /*mem_size   =*/ mem_size,
836            /*mem_buffer =*/ nullptr,
837            /*no_alloc   =*/ params.no_alloc,
838        };
839
840        *params.ctx = ggml_init(pdata);
841        if (*params.ctx == nullptr) {
842            GGML_LOG_ERROR("%s: failed to initialize ggml context for storing tensors\n", __func__);
843            gguf_free(ctx);
844            return nullptr;
845        }
846
847        struct ggml_context * ctx_data = *params.ctx;
848
849        struct ggml_tensor * data = nullptr;
850
851        if (!params.no_alloc) {
852            data = ggml_new_tensor_1d(ctx_data, GGML_TYPE_I8, ctx->size);
853
854            ok = ok && data != nullptr;
855
856            if (ok) {
857                ggml_set_name(data, "GGUF tensor data binary blob");
858            }
859
860            // read the binary blob with the tensor data
861            ok = ok && gr.read(data->data, ctx->size);
862
863            if (!ok) {
864                GGML_LOG_ERROR("%s: failed to read tensor data binary blob\n", __func__);
865                ggml_free(ctx_data);
866                *params.ctx = nullptr;
867                gguf_free(ctx);
868                return nullptr;
869            }
870
871            ctx->data = data->data;
872        }
873
874        ggml_set_no_alloc(ctx_data, true);
875
876        // create the tensors
877        for (size_t i = 0; i < ctx->info.size(); ++i) {
878            const struct gguf_tensor_info & info = ctx->info[i];
879
880            struct ggml_tensor * cur = ggml_new_tensor(ctx_data, info.t.type, GGML_MAX_DIMS, info.t.ne);
881
882            ok = ok && cur != nullptr;
883
884            if (!ok) {
885                break;
886            }
887
888            ggml_set_name(cur, info.t.name);
889
890            // point the data member to the appropriate location in the binary blob using the tensor info
891            if (!params.no_alloc) {
892                cur->data = (char *) data->data + info.offset;
893            }
894        }
895
896        if (!ok) {
897            GGML_LOG_ERROR("%s: failed to create tensors\n", __func__);
898            ggml_free(ctx_data);
899            *params.ctx = nullptr;
900            gguf_free(ctx);
901            return nullptr;
902        }
903
904        ggml_set_no_alloc(ctx_data, params.no_alloc);
905    }
906
907    return ctx;
908}
909
910struct gguf_context * gguf_init_from_callback(gguf_reader_callback_t callback, void * userdata, size_t max_chunk_read, uint64_t max_expected_size, struct gguf_init_params params) {
911    if (callback == nullptr) {
912        return nullptr;
913    }
914
915    const struct gguf_reader gr(callback, userdata, max_chunk_read == 0 ? SIZE_MAX : max_chunk_read, 0, max_expected_size);
916    return gguf_init_from_reader(gr, params);
917}
918
919struct gguf_file_reader {
920    FILE * file;
921    uint64_t offset;
922};
923
924static size_t gguf_file_reader_callback(void * userdata, void * output, uint64_t offset, size_t len) {
925    GGML_ASSERT(len > 0);
926
927    gguf_file_reader & reader = *static_cast<gguf_file_reader *>(userdata);
928
929    if (reader.offset != offset) {
930        if (offset > INT64_MAX || gguf_fseek(reader.file, static_cast<int64_t>(offset), SEEK_SET) != 0) {
931            return 0;
932        }
933
934        reader.offset = offset;
935    }
936
937    const size_t nread = fread(static_cast<uint8_t *>(output), 1, len, reader.file);
938    reader.offset += nread;
939    return nread;
940}
941
942struct gguf_context * gguf_init_from_file_ptr(FILE * file, struct gguf_init_params params) {
943    if (!file) {
944        return nullptr;
945    }
946
947    const int64_t cur = gguf_ftell(file);
948    if (cur < 0) {
949        return nullptr;
950    }
951
952    gguf_file_reader reader = {
953        /*.file   = */ file,
954        /*.offset = */ static_cast<uint64_t>(cur),
955    };
956    const struct gguf_reader gr(gguf_file_reader_callback, &reader, SIZE_MAX, reader.offset, gguf_reader::file_remain(file));
957    return gguf_init_from_reader(gr, params);
958}
959
960struct gguf_buffer_reader {
961    const uint8_t * data;
962    size_t          size;
963};
964
965static size_t gguf_buffer_reader_callback(void * userdata, void * output, uint64_t offset, size_t len) {
966    GGML_ASSERT(len > 0);
967
968    const gguf_buffer_reader & reader = *static_cast<gguf_buffer_reader *>(userdata);
969
970    if (offset > reader.size || len > reader.size - offset) {
971        return 0;
972    }
973
974    const size_t data_offset = static_cast<size_t>(offset);
975    const size_t nread = std::min(len, reader.size - data_offset);
976    memcpy(static_cast<uint8_t *>(output), reader.data + data_offset, nread);
977    return nread;
978}
979
980struct gguf_context * gguf_init_from_buffer(const void * data, size_t size, struct gguf_init_params params) {
981    if (data == nullptr || size == 0) {
982        return nullptr;
983    }
984
985    gguf_buffer_reader reader = {
986        /*.data = */ static_cast<const uint8_t *>(data),
987        /*.size = */ size,
988    };
989    const struct gguf_reader gr(gguf_buffer_reader_callback, &reader, SIZE_MAX, 0, size);
990    return gguf_init_from_reader(gr, params);
991}
992
993struct gguf_context * gguf_init_from_file(const char * fname, struct gguf_init_params params) {
994    FILE * file = ggml_fopen(fname, "rb");
995
996    if (!file) {
997        GGML_LOG_ERROR("%s: failed to open GGUF file '%s' (%s)\n", __func__, fname, strerror(errno));
998        return nullptr;
999    }
1000
1001    struct gguf_context * result = gguf_init_from_file_ptr(file, params);
1002    fclose(file);
1003    return result;
1004}
1005
1006void gguf_free(struct gguf_context * ctx) {
1007    if (ctx == nullptr) {
1008        return;
1009    }
1010    delete ctx;
1011}
1012
1013const char * gguf_type_name(enum gguf_type type) {
1014    auto it = GGUF_TYPE_NAME.find(type);
1015    return it == GGUF_TYPE_NAME.end() ? nullptr : it->second;
1016}
1017
1018uint32_t gguf_get_version(const struct gguf_context * ctx) {
1019    return ctx->version;
1020}
1021
1022size_t gguf_get_alignment(const struct gguf_context * ctx) {
1023    return ctx->alignment;
1024}
1025
1026size_t gguf_get_data_offset(const struct gguf_context * ctx) {
1027    return ctx->offset;
1028}
1029
1030int64_t gguf_get_n_kv(const struct gguf_context * ctx) {
1031    return ctx->kv.size();
1032}
1033
1034int64_t gguf_find_key(const struct gguf_context * ctx, const char * key) {
1035    // return -1 if key not found
1036    int64_t keyfound = -1;
1037
1038    const int64_t n_kv = gguf_get_n_kv(ctx);
1039
1040    for (int64_t i = 0; i < n_kv; ++i) {
1041        if (strcmp(key, gguf_get_key(ctx, i)) == 0) {
1042            keyfound = i;
1043            break;
1044        }
1045    }
1046
1047    return keyfound;
1048}
1049
1050const char * gguf_get_key(const struct gguf_context * ctx, int64_t key_id) {
1051    GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx));
1052    return ctx->kv[key_id].get_key().c_str();
1053}
1054
1055enum gguf_type gguf_get_kv_type(const struct gguf_context * ctx, int64_t key_id) {
1056    GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx));
1057    return ctx->kv[key_id].is_array ? GGUF_TYPE_ARRAY : ctx->kv[key_id].get_type();
1058}
1059
1060enum gguf_type gguf_get_arr_type(const struct gguf_context * ctx, int64_t key_id) {
1061    GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx));
1062    GGML_ASSERT(ctx->kv[key_id].is_array);
1063    return ctx->kv[key_id].get_type();
1064}
1065
1066const void * gguf_get_arr_data(const struct gguf_context * ctx, int64_t key_id) {
1067    GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx));
1068    GGML_ASSERT(ctx->kv[key_id].get_type() != GGUF_TYPE_STRING);
1069    return ctx->kv[key_id].data.data();
1070}
1071
1072const char * gguf_get_arr_str(const struct gguf_context * ctx, int64_t key_id, size_t i) {
1073    GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx));
1074    GGML_ASSERT(ctx->kv[key_id].get_type() == GGUF_TYPE_STRING);
1075    return ctx->kv[key_id].data_string[i].c_str();
1076}
1077
1078size_t gguf_get_arr_n(const struct gguf_context * ctx, int64_t key_id) {
1079    GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx));
1080
1081    if (ctx->kv[key_id].type == GGUF_TYPE_STRING) {
1082        return ctx->kv[key_id].data_string.size();
1083    }
1084
1085    const size_t type_size = gguf_type_size(ctx->kv[key_id].type);
1086    GGML_ASSERT(ctx->kv[key_id].data.size() % type_size == 0);
1087    return ctx->kv[key_id].data.size() / type_size;
1088}
1089
1090uint8_t gguf_get_val_u8(const struct gguf_context * ctx, int64_t key_id) {
1091    GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx));
1092    GGML_ASSERT(ctx->kv[key_id].get_ne() == 1);
1093    return ctx->kv[key_id].get_val<uint8_t>();
1094}
1095
1096int8_t gguf_get_val_i8(const struct gguf_context * ctx, int64_t key_id) {
1097    GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx));
1098    GGML_ASSERT(ctx->kv[key_id].get_ne() == 1);
1099    return ctx->kv[key_id].get_val<int8_t>();
1100}
1101
1102uint16_t gguf_get_val_u16(const struct gguf_context * ctx, int64_t key_id) {
1103    GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx));
1104    GGML_ASSERT(ctx->kv[key_id].get_ne() == 1);
1105    return ctx->kv[key_id].get_val<uint16_t>();
1106}
1107
1108int16_t gguf_get_val_i16(const struct gguf_context * ctx, int64_t key_id) {
1109    GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx));
1110    GGML_ASSERT(ctx->kv[key_id].get_ne() == 1);
1111    return ctx->kv[key_id].get_val<int16_t>();
1112}
1113
1114uint32_t gguf_get_val_u32(const struct gguf_context * ctx, int64_t key_id) {
1115    GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx));
1116    GGML_ASSERT(ctx->kv[key_id].get_ne() == 1);
1117    return ctx->kv[key_id].get_val<uint32_t>();
1118}
1119
1120int32_t gguf_get_val_i32(const struct gguf_context * ctx, int64_t key_id) {
1121    GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx));
1122    GGML_ASSERT(ctx->kv[key_id].get_ne() == 1);
1123    return ctx->kv[key_id].get_val<int32_t>();
1124}
1125
1126float gguf_get_val_f32(const struct gguf_context * ctx, int64_t key_id) {
1127    GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx));
1128    GGML_ASSERT(ctx->kv[key_id].get_ne() == 1);
1129    return ctx->kv[key_id].get_val<float>();
1130}
1131
1132uint64_t gguf_get_val_u64(const struct gguf_context * ctx, int64_t key_id) {
1133    GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx));
1134    GGML_ASSERT(ctx->kv[key_id].get_ne() == 1);
1135    return ctx->kv[key_id].get_val<uint64_t>();
1136}
1137
1138int64_t gguf_get_val_i64(const struct gguf_context * ctx, int64_t key_id) {
1139    GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx));
1140    GGML_ASSERT(ctx->kv[key_id].get_ne() == 1);
1141    return ctx->kv[key_id].get_val<int64_t>();
1142}
1143
1144double gguf_get_val_f64(const struct gguf_context * ctx, int64_t key_id) {
1145    GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx));
1146    GGML_ASSERT(ctx->kv[key_id].get_ne() == 1);
1147    return ctx->kv[key_id].get_val<double>();
1148}
1149
1150bool gguf_get_val_bool(const struct gguf_context * ctx, int64_t key_id) {
1151    GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx));
1152    GGML_ASSERT(ctx->kv[key_id].get_ne() == 1);
1153    return ctx->kv[key_id].get_val<bool>();
1154}
1155
1156const char * gguf_get_val_str(const struct gguf_context * ctx, int64_t key_id) {
1157    GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx));
1158    GGML_ASSERT(ctx->kv[key_id].get_ne() == 1);
1159    return ctx->kv[key_id].get_val<std::string>().c_str();
1160}
1161
1162const void * gguf_get_val_data(const struct gguf_context * ctx, int64_t key_id) {
1163    GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx));
1164    GGML_ASSERT(ctx->kv[key_id].get_ne() == 1);
1165    GGML_ASSERT(ctx->kv[key_id].get_type() != GGUF_TYPE_STRING);
1166    return ctx->kv[key_id].data.data();
1167}
1168
1169int64_t gguf_get_n_tensors(const struct gguf_context * ctx) {
1170    return ctx->info.size();
1171}
1172
1173int64_t gguf_find_tensor(const struct gguf_context * ctx, const char * name) {
1174    // return -1 if tensor not found
1175    int64_t tensor_id = -1;
1176
1177    const int64_t n_tensors = gguf_get_n_tensors(ctx);
1178
1179    for (int64_t i = 0; i < n_tensors; ++i) {
1180        if (strcmp(name, gguf_get_tensor_name(ctx, i)) == 0) {
1181            tensor_id = i;
1182            break;
1183        }
1184    }
1185
1186    return tensor_id;
1187}
1188
1189size_t gguf_get_tensor_offset(const struct gguf_context * ctx, int64_t tensor_id) {
1190    GGML_ASSERT(tensor_id >= 0 && tensor_id < gguf_get_n_tensors(ctx));
1191    return ctx->info[tensor_id].offset;
1192}
1193
1194const char * gguf_get_tensor_name(const struct gguf_context * ctx, int64_t tensor_id) {
1195    GGML_ASSERT(tensor_id >= 0 && tensor_id < gguf_get_n_tensors(ctx));
1196    return ctx->info[tensor_id].t.name;
1197}
1198
1199const int64_t * gguf_get_tensor_ne(const struct gguf_context * ctx, int64_t tensor_id) {
1200    GGML_ASSERT(tensor_id >= 0 && tensor_id < gguf_get_n_tensors(ctx));
1201    return ctx->info[tensor_id].t.ne;
1202}
1203
1204enum ggml_type gguf_get_tensor_type(const struct gguf_context * ctx, int64_t tensor_id) {
1205    GGML_ASSERT(tensor_id >= 0 && tensor_id < gguf_get_n_tensors(ctx));
1206    return ctx->info[tensor_id].t.type;
1207}
1208
1209size_t gguf_get_tensor_size(const struct gguf_context * ctx, int64_t tensor_id) {
1210    GGML_ASSERT(tensor_id >= 0 && tensor_id < gguf_get_n_tensors(ctx));
1211    return ggml_nbytes(&ctx->info[tensor_id].t);
1212}
1213
1214int64_t gguf_remove_key(struct gguf_context * ctx, const char * key) {
1215    const int64_t key_id = gguf_find_key(ctx, key);
1216    if (key_id >= 0) {
1217        ctx->kv.erase(ctx->kv.begin() + key_id);
1218    }
1219    return key_id;
1220}
1221
1222template<typename T>
1223static void gguf_check_reserved_keys(const std::string & key, const T val) {
1224    if (key == GGUF_KEY_GENERAL_ALIGNMENT) {
1225        if constexpr (std::is_same<T, uint32_t>::value) {
1226            GGML_ASSERT(val > 0 && (val & (val - 1)) == 0 && GGUF_KEY_GENERAL_ALIGNMENT " must be power of 2");
1227        } else {
1228            GGML_UNUSED(val);
1229            GGML_ABORT(GGUF_KEY_GENERAL_ALIGNMENT " must be type u32");
1230        }
1231    }
1232}
1233
1234void gguf_set_val_u8(struct gguf_context * ctx, const char * key, uint8_t val) {
1235    gguf_check_reserved_keys(key, val);
1236    gguf_remove_key(ctx, key);
1237    ctx->kv.emplace_back(key, val);
1238}
1239
1240void gguf_set_val_i8(struct gguf_context * ctx, const char * key, int8_t val) {
1241    gguf_check_reserved_keys(key, val);
1242    gguf_remove_key(ctx, key);
1243    ctx->kv.emplace_back(key, val);
1244}
1245
1246void gguf_set_val_u16(struct gguf_context * ctx, const char * key, uint16_t val) {
1247    gguf_check_reserved_keys(key, val);
1248    gguf_remove_key(ctx, key);
1249    ctx->kv.emplace_back(key, val);
1250}
1251
1252void gguf_set_val_i16(struct gguf_context * ctx, const char * key, int16_t val) {
1253    gguf_check_reserved_keys(key, val);
1254    gguf_remove_key(ctx, key);
1255    ctx->kv.emplace_back(key, val);
1256}
1257
1258void gguf_set_val_u32(struct gguf_context * ctx, const char * key, uint32_t val) {
1259    gguf_check_reserved_keys(key, val);
1260    gguf_remove_key(ctx, key);
1261    ctx->kv.emplace_back(key, val);
1262}
1263
1264void gguf_set_val_i32(struct gguf_context * ctx, const char * key, int32_t val) {
1265    gguf_check_reserved_keys(key, val);
1266    gguf_remove_key(ctx, key);
1267    ctx->kv.emplace_back(key, val);
1268}
1269
1270void gguf_set_val_f32(struct gguf_context * ctx, const char * key, float val) {
1271    gguf_check_reserved_keys(key, val);
1272    gguf_remove_key(ctx, key);
1273    ctx->kv.emplace_back(key, val);
1274}
1275
1276void gguf_set_val_u64(struct gguf_context * ctx, const char * key, uint64_t val) {
1277    gguf_check_reserved_keys(key, val);
1278    gguf_remove_key(ctx, key);
1279    ctx->kv.emplace_back(key, val);
1280}
1281
1282void gguf_set_val_i64(struct gguf_context * ctx, const char * key, int64_t val) {
1283    gguf_check_reserved_keys(key, val);
1284    gguf_remove_key(ctx, key);
1285    ctx->kv.emplace_back(key, val);
1286}
1287
1288void gguf_set_val_f64(struct gguf_context * ctx, const char * key, double val) {
1289    gguf_check_reserved_keys(key, val);
1290    gguf_remove_key(ctx, key);
1291    ctx->kv.emplace_back(key, val);
1292}
1293
1294void gguf_set_val_bool(struct gguf_context * ctx, const char * key, bool val) {
1295    gguf_check_reserved_keys(key, val);
1296    gguf_remove_key(ctx, key);
1297    ctx->kv.emplace_back(key, val);
1298}
1299
1300void gguf_set_val_str(struct gguf_context * ctx, const char * key, const char * val) {
1301    gguf_check_reserved_keys(key, val);
1302    gguf_remove_key(ctx, key);
1303    ctx->kv.emplace_back(key, std::string(val));
1304}
1305
1306void gguf_set_arr_data(struct gguf_context * ctx, const char * key, enum gguf_type type, const void * data, size_t n) {
1307    gguf_check_reserved_keys(key, data);
1308    gguf_remove_key(ctx, key);
1309
1310    const size_t nbytes = n*gguf_type_size(type);
1311    std::vector<int8_t> tmp(nbytes);
1312    if (!tmp.empty()) {
1313        memcpy(tmp.data(), data, nbytes);
1314    }
1315    ctx->kv.emplace_back(key, tmp);
1316    ctx->kv.back().cast(type);
1317}
1318
1319void gguf_set_arr_str(struct gguf_context * ctx, const char * key, const char ** data, size_t n) {
1320    gguf_check_reserved_keys(key, data);
1321    gguf_remove_key(ctx, key);
1322
1323    std::vector<std::string> tmp(n);
1324    for (size_t i = 0; i < n; ++i) {
1325        tmp[i] = data[i];
1326    }
1327    ctx->kv.emplace_back(key, tmp);
1328}
1329
1330// set or add KV pairs from another context
1331void gguf_set_kv(struct gguf_context * ctx, const struct gguf_context * src) {
1332    const int64_t n_kv = gguf_get_n_kv(src);
1333    for (int64_t i = 0; i < n_kv; ++i) {
1334        const struct gguf_kv & kv = src->kv[i];
1335
1336        if (!kv.is_array) {
1337            switch (kv.get_type()) {
1338                case GGUF_TYPE_UINT8:   gguf_set_val_u8  (ctx, kv.get_key().c_str(), kv.get_val<uint8_t>());             break;
1339                case GGUF_TYPE_INT8:    gguf_set_val_i8  (ctx, kv.get_key().c_str(), kv.get_val<int8_t>());              break;
1340                case GGUF_TYPE_UINT16:  gguf_set_val_u16 (ctx, kv.get_key().c_str(), kv.get_val<uint16_t>());            break;
1341                case GGUF_TYPE_INT16:   gguf_set_val_i16 (ctx, kv.get_key().c_str(), kv.get_val<int16_t>());             break;
1342                case GGUF_TYPE_UINT32:  gguf_set_val_u32 (ctx, kv.get_key().c_str(), kv.get_val<uint32_t>());            break;
1343                case GGUF_TYPE_INT32:   gguf_set_val_i32 (ctx, kv.get_key().c_str(), kv.get_val<int32_t>());             break;
1344                case GGUF_TYPE_FLOAT32: gguf_set_val_f32 (ctx, kv.get_key().c_str(), kv.get_val<float>());               break;
1345                case GGUF_TYPE_UINT64:  gguf_set_val_u64 (ctx, kv.get_key().c_str(), kv.get_val<uint64_t>());            break;
1346                case GGUF_TYPE_INT64:   gguf_set_val_i64 (ctx, kv.get_key().c_str(), kv.get_val<int64_t>());             break;
1347                case GGUF_TYPE_FLOAT64: gguf_set_val_f64 (ctx, kv.get_key().c_str(), kv.get_val<double>());              break;
1348                case GGUF_TYPE_BOOL:    gguf_set_val_bool(ctx, kv.get_key().c_str(), kv.get_val<bool>());                break;
1349                case GGUF_TYPE_STRING:  gguf_set_val_str (ctx, kv.get_key().c_str(), kv.get_val<std::string>().c_str()); break;
1350                case GGUF_TYPE_ARRAY:
1351                default: GGML_ABORT("invalid type");
1352            }
1353            continue;
1354        }
1355
1356        const size_t ne = kv.get_ne();
1357
1358        switch (kv.get_type()) {
1359            case GGUF_TYPE_UINT8:
1360            case GGUF_TYPE_INT8:
1361            case GGUF_TYPE_UINT16:
1362            case GGUF_TYPE_INT16:
1363            case GGUF_TYPE_UINT32:
1364            case GGUF_TYPE_INT32:
1365            case GGUF_TYPE_FLOAT32:
1366            case GGUF_TYPE_UINT64:
1367            case GGUF_TYPE_INT64:
1368            case GGUF_TYPE_FLOAT64:
1369            case GGUF_TYPE_BOOL: {
1370                gguf_set_arr_data(ctx, kv.get_key().c_str(), kv.get_type(), kv.data.data(), ne);
1371            } break;
1372            case GGUF_TYPE_STRING: {
1373                std::vector<const char *> tmp(ne);
1374                for (size_t j = 0; j < ne; ++j) {
1375                    tmp[j] = kv.data_string[j].c_str();
1376                }
1377                gguf_set_arr_str(ctx, kv.get_key().c_str(), tmp.data(), ne);
1378            } break;
1379            case GGUF_TYPE_ARRAY:
1380            default: GGML_ABORT("invalid type");
1381        }
1382    }
1383}
1384
1385void gguf_add_tensor(
1386             struct gguf_context * ctx,
1387        const struct ggml_tensor * tensor) {
1388    GGML_ASSERT(tensor);
1389    if (gguf_find_tensor(ctx, tensor->name) != -1) {
1390        GGML_ABORT("duplicate tensor name: %s", tensor->name);
1391    }
1392
1393    struct gguf_tensor_info ti;
1394    ti.t = *tensor;
1395    ti.offset = ctx->info.empty() ? 0 :
1396        ctx->info.back().offset + GGML_PAD(ggml_nbytes(&ctx->info.back().t), ctx->alignment);
1397    ctx->info.push_back(ti);
1398}
1399
1400void gguf_set_tensor_type(struct gguf_context * ctx, const char * name, enum ggml_type type) {
1401    const int64_t tensor_id = gguf_find_tensor(ctx, name);
1402    if (tensor_id < 0) {
1403        GGML_ABORT("tensor not found: %s", name);
1404    }
1405    struct ggml_tensor * tensor = &ctx->info[tensor_id].t;
1406    const size_t  type_size = ggml_type_size(type);
1407    const int64_t blck_size = ggml_blck_size(type);
1408
1409    tensor->type = type;
1410    GGML_ASSERT(tensor->ne[0] % blck_size == 0 && "tensor row size not divisible by block size of new type");
1411
1412    tensor->nb[0] = type_size;
1413    tensor->nb[1] = tensor->nb[0]*(tensor->ne[0]/blck_size);
1414    for (int i = 2; i < GGML_MAX_DIMS; i++) {
1415        tensor->nb[i] = tensor->nb[i - 1]*tensor->ne[i - 1];
1416    }
1417
1418    // update offsets
1419    const int64_t n_tensors = gguf_get_n_tensors(ctx);
1420    for (int64_t i = tensor_id + 1; i < n_tensors; ++i) {
1421        ctx->info[i].offset = ctx->info[i - 1].offset + GGML_PAD(ggml_nbytes(&ctx->info[i - 1].t), ctx->alignment);
1422    }
1423}
1424
1425void gguf_set_tensor_data(struct gguf_context * ctx, const char * name, const void * data) {
1426    const int64_t tensor_id = gguf_find_tensor(ctx, name);
1427    if (tensor_id < 0) {
1428        GGML_ABORT("tensor not found: %s", name);
1429    }
1430
1431    ctx->info[tensor_id].t.data = (void *)(uintptr_t)data; // double cast suppresses warning about casting away const
1432}
1433
1434struct gguf_writer_base {
1435    size_t written_bytes {0u};
1436
1437    virtual ~gguf_writer_base(void) = default;
1438
1439    // we bet on devirtualization
1440    virtual void write(int8_t val) = 0;
1441    virtual void write(const std::vector<int8_t> & val) = 0;
1442    virtual void write_tensor_data(const struct gguf_tensor_info & info, size_t offset_data, size_t alignment) = 0;
1443
1444    template <typename T>
1445    void write(const T & val) {
1446        for (size_t i = 0; i < sizeof(val); ++i) {
1447            write(reinterpret_cast<const int8_t *>(&val)[i]);
1448        }
1449    }
1450
1451    void write(const bool & val) {
1452        const int8_t val8 = val ? 1 : 0;
1453        write(val8);
1454    }
1455
1456    void write(const std::string & val) {
1457        {
1458            const uint64_t n = val.length();
1459            write(n);
1460        }
1461        for (size_t i = 0; i < val.length(); ++i) {
1462            write((val.data())[i]);
1463        }
1464    }
1465
1466    void write(const char * val) {
1467        write(std::string(val));
1468    }
1469
1470    void write(const enum ggml_type & val) {
1471        write(int32_t(val));
1472    }
1473
1474    void write(const enum gguf_type & val) {
1475        write(int32_t(val));
1476    }
1477
1478    void write(const struct gguf_kv & kv) {
1479        const uint64_t ne = kv.get_ne();
1480
1481        write(kv.get_key());
1482
1483        if (kv.is_array) {
1484            write(GGUF_TYPE_ARRAY);
1485            write(kv.get_type());
1486            write(ne);
1487        } else {
1488            write(kv.get_type());
1489        }
1490
1491        switch (kv.get_type()) {
1492            case GGUF_TYPE_UINT8:
1493            case GGUF_TYPE_INT8:
1494            case GGUF_TYPE_UINT16:
1495            case GGUF_TYPE_INT16:
1496            case GGUF_TYPE_UINT32:
1497            case GGUF_TYPE_INT32:
1498            case GGUF_TYPE_FLOAT32:
1499            case GGUF_TYPE_UINT64:
1500            case GGUF_TYPE_INT64:
1501            case GGUF_TYPE_FLOAT64: {
1502                write(kv.data);
1503            } break;
1504            case GGUF_TYPE_BOOL: {
1505                for (size_t i = 0; i < ne; ++i) {
1506                    write(kv.get_val<bool>(i));
1507                }
1508            } break;
1509            case GGUF_TYPE_STRING: {
1510                for (size_t i = 0; i < ne; ++i) {
1511                    write(kv.get_val<std::string>(i));
1512                }
1513            } break;
1514            case GGUF_TYPE_ARRAY:
1515            default: GGML_ABORT("invalid type");
1516        }
1517    }
1518
1519    void write_tensor_meta(const struct gguf_tensor_info & info) {
1520        write(info.t.name);
1521
1522        const uint32_t n_dims = ggml_n_dims(&info.t);
1523        write(n_dims);
1524
1525        for (uint32_t j = 0; j < n_dims; ++j) {
1526            write(info.t.ne[j]);
1527        }
1528        write(info.t.type);
1529        write(info.offset);
1530    }
1531
1532    void pad(const size_t alignment) {
1533        while (written_bytes % alignment != 0) {
1534            const int8_t zero = 0;
1535            write(zero);
1536        }
1537    }
1538};
1539
1540// vector buffer based writer
1541struct gguf_writer_buf final : public gguf_writer_base {
1542    std::vector<int8_t> & buf;
1543
1544    gguf_writer_buf(std::vector<int8_t> & buf) : buf(buf) {}
1545
1546    using gguf_writer_base::write;
1547
1548    void write(const int8_t val) override {
1549        buf.push_back(val);
1550        written_bytes++;
1551    }
1552
1553    void write(const std::vector<int8_t> & val) override {
1554        buf.insert(buf.end(), val.begin(), val.end());
1555        written_bytes += val.size();
1556    }
1557
1558    void write_tensor_data(const struct gguf_tensor_info & info, const size_t offset_data, const size_t alignment) override {
1559        GGML_ASSERT(buf.size() - offset_data == info.offset);
1560
1561        GGML_ASSERT(ggml_is_contiguous(&info.t));
1562        const size_t offset = buf.size();
1563        const size_t nbytes = ggml_nbytes(&info.t);
1564
1565        buf.resize(offset + nbytes);
1566        if (info.t.buffer) {
1567            ggml_backend_tensor_get(&info.t, buf.data() + offset, 0, nbytes);
1568        } else {
1569            GGML_ASSERT(info.t.data);
1570            memcpy(buf.data() + offset, info.t.data, nbytes);
1571        }
1572        written_bytes += nbytes;
1573
1574        pad(alignment);
1575    }
1576};
1577
1578// file based writer
1579struct gguf_writer_file final : public gguf_writer_base {
1580    FILE * file;
1581
1582    gguf_writer_file(FILE* file) : file(file) {}
1583
1584    using gguf_writer_base::write;
1585
1586    void write(const int8_t val) override {
1587        const auto real_val = static_cast<uint8_t>(val);
1588        const auto ret = fputc(real_val, file);
1589        written_bytes++;
1590        if (ret != real_val) {
1591            throw std::runtime_error("unexpected fputc result '" + std::to_string(ret) + "' instead of '" + std::to_string((int)real_val) + "'");
1592        }
1593    }
1594
1595    void write(const std::vector<int8_t> & val) override {
1596        const auto ret = fwrite(val.data(), 1, val.size(), file);
1597        written_bytes += val.size();
1598        if (ret != val.size()) {
1599            throw std::runtime_error("unexpected fwrite number of bytes written, '" + std::to_string(ret) + "' instead of '" + std::to_string(val.size()) + "'");
1600        }
1601    }
1602
1603    void write_tensor_data(const struct gguf_tensor_info & info, const size_t offset_data, const size_t alignment) override {
1604        GGML_ASSERT(written_bytes - offset_data == info.offset);
1605
1606        GGML_ASSERT(ggml_is_contiguous(&info.t));
1607        const size_t nbytes = ggml_nbytes(&info.t);
1608
1609        std::vector<int8_t> buf(nbytes);
1610        if (info.t.buffer) {
1611            ggml_backend_tensor_get(&info.t, buf.data(), 0, nbytes);
1612        } else {
1613            GGML_ASSERT(info.t.data);
1614            memcpy(buf.data(), info.t.data, nbytes);
1615        }
1616        write(buf);
1617
1618        pad(alignment);
1619    }
1620};
1621
1622template <typename writer_t>
1623static void gguf_write_out(const struct gguf_context * ctx, writer_t & gw, bool only_meta) {
1624    const int64_t n_kv      = gguf_get_n_kv(ctx);
1625    const int64_t n_tensors = gguf_get_n_tensors(ctx);
1626
1627    // write header
1628    gw.write(GGUF_MAGIC[0]);
1629    gw.write(GGUF_MAGIC[1]);
1630    gw.write(GGUF_MAGIC[2]);
1631    gw.write(GGUF_MAGIC[3]);
1632    gw.write(ctx->version);
1633    gw.write(n_tensors);
1634    gw.write(n_kv);
1635
1636    // write key-value pairs
1637    for (int64_t i = 0; i < n_kv; ++i) {
1638        gw.write(ctx->kv[i]);
1639    }
1640
1641    // write tensor info
1642    for (int64_t i = 0; i < n_tensors; ++i) {
1643        gw.write_tensor_meta(ctx->info[i]);
1644    }
1645
1646    // we require the data section to be aligned
1647    gw.pad(ctx->alignment);
1648
1649    if (only_meta) {
1650        return;
1651    }
1652
1653    const size_t offset_data = gw.written_bytes;
1654
1655    // write tensor data
1656    for (int64_t i = 0; i < n_tensors; ++i) {
1657        gw.write_tensor_data(ctx->info[i], offset_data, ctx->alignment);
1658    }
1659}
1660
1661void gguf_write_to_buf(const struct gguf_context * ctx, std::vector<int8_t> & buf, bool only_meta) {
1662    gguf_writer_buf gw(buf);
1663    gguf_write_out(ctx, gw, only_meta);
1664}
1665
1666bool gguf_write_to_file_ptr(const struct gguf_context * ctx, FILE * file, bool only_meta) {
1667    GGML_ASSERT(file);
1668
1669    try {
1670        gguf_writer_file gw(file);
1671        gguf_write_out(ctx, gw, only_meta);
1672    } catch (const std::runtime_error& ex) {
1673        GGML_LOG_ERROR("%s: failed to write GGUF data: %s\n", __func__, ex.what());
1674        return false;
1675    }
1676    return true;
1677}
1678
1679bool gguf_write_to_file(const struct gguf_context * ctx, const char * fname, bool only_meta) {
1680    FILE * file = ggml_fopen(fname, "wb");
1681
1682    if (!file) {
1683        GGML_LOG_ERROR("%s: failed to open file '%s' for writing GGUF data\n", __func__, fname);
1684        return false;
1685    }
1686
1687    const bool success = gguf_write_to_file_ptr(ctx, file, only_meta);
1688    if (!success) {
1689        GGML_LOG_ERROR("%s: failed to write GGUF data into '%s'\n", __func__, fname);
1690    }
1691
1692    fclose(file);
1693    return success;
1694}
1695
1696size_t gguf_get_meta_size(const struct gguf_context * ctx) {
1697    // only return size
1698    std::vector<int8_t> buf;
1699    gguf_write_to_buf(ctx, buf, /*only_meta =*/ true);
1700    return buf.size();
1701}
1702
1703void gguf_get_meta_data(const struct gguf_context * ctx, void * data) {
1704    std::vector<int8_t> buf;
1705    gguf_write_to_buf(ctx, buf, /*only_meta =*/ true);
1706    memcpy(data, buf.data(), buf.size());
1707}