API to remap/rewind + mapping in write mode.

This commit is contained in:
antirez
2023-12-25 10:45:38 +01:00
parent f400e8a36f
commit 3eb30c1872
2 changed files with 46 additions and 23 deletions

View File

@@ -93,49 +93,70 @@ uint64_t gguf_value_len(uint32_t type, union gguf_value *val) {
/* Open a GGUF file and return a parsing context. */ /* Open a GGUF file and return a parsing context. */
gguf_ctx *gguf_init(char *filename) { gguf_ctx *gguf_init(char *filename) {
struct stat sb; int fd = open(filename,O_RDWR);
int fd = open(filename,O_RDONLY);
if (fd == -1) return NULL; if (fd == -1) return NULL;
if (fstat(fd,&sb) == -1) {
close(fd);
return NULL;
}
/* Now that we have an open file and its total size, let's /* Mapping successful. We can create our context object. */
* mmap it. */ gguf_ctx *ctx = malloc(sizeof(*ctx));
void *mapped = mmap(0,sb.st_size,PROT_READ,MAP_PRIVATE,fd,0); memset(ctx,0,sizeof(*ctx));
if (mapped == MAP_FAILED) { ctx->fd = fd;
close(fd); ctx->alignment = 32; // Default alighment of GGUF files.
ctx->data_off = 0; // Set later.
if (gguf_remap(ctx) == 0) {
gguf_end(ctx);
return NULL; return NULL;
} }
gguf_rewind(ctx);
return ctx;
}
/* Set the context to read the first key-value entry in the GGUF
* file and then all the rest. Is used when creating a new context
* and also when you want to restart scanning the key-value
* items in the file. */
void gguf_rewind(gguf_ctx *ctx) {
ctx->off = sizeof(struct gguf_header);
ctx->left_kv = ctx->header->metadata_kv_count;
ctx->left_tensors = ctx->header->tensor_count;
}
/* map or re-map the GGUF file inside the context pointers to
* header and data, also calculating the file length. This is
* used when creating a context, but also after the user write
* to the file extending it, and requires to view again the
* whole updated file.
*
* Return 1 on success, 0 on error. */
int gguf_remap(gguf_ctx *ctx) {
struct stat sb;
/* Unmap if the file was already memory mapped. */
if (ctx->data) munmap(ctx->data,ctx->size);
/* Get the size of the file to map, then map it. */
if (fstat(ctx->fd,&sb) == -1) return 0;
void *mapped = mmap(0,sb.st_size,PROT_READ|PROT_WRITE,MAP_SHARED,ctx->fd,0);
if (mapped == MAP_FAILED) return 0;
/* Minimal sanity check... */ /* Minimal sanity check... */
if (sb.st_size < (signed)sizeof(struct gguf_header) || if (sb.st_size < (signed)sizeof(struct gguf_header) ||
memcmp(mapped,"GGUF",4) != 0) memcmp(mapped,"GGUF",4) != 0)
{ {
errno = EINVAL; errno = EINVAL;
return NULL; return 0;
} }
/* Mapping successful. We can create our context object. */
gguf_ctx *ctx = malloc(sizeof(*ctx));
ctx->fd = fd;
ctx->data = mapped; ctx->data = mapped;
ctx->header = mapped; ctx->header = mapped;
ctx->size = sb.st_size; ctx->size = sb.st_size;
ctx->off = sizeof(struct gguf_header); return 1;
ctx->left_kv = ctx->header->metadata_kv_count;
ctx->left_tensors = ctx->header->tensor_count;
ctx->alignment = 32; // Default alighment of GGUF files.
ctx->data_off = 0; // Set later.
return ctx;
} }
/* Cleanup needed after gguf_init(), to terminate the context /* Cleanup needed after gguf_init(), to terminate the context
* and cleanup resources. */ * and cleanup resources. */
void gguf_end(gguf_ctx *ctx) { void gguf_end(gguf_ctx *ctx) {
if (ctx == NULL) return; if (ctx == NULL) return;
munmap(ctx->data,ctx->size); if (ctx->data) munmap(ctx->data,ctx->size);
close(ctx->fd); close(ctx->fd);
free(ctx); free(ctx);
} }

View File

@@ -159,6 +159,8 @@ typedef struct {
/* =============================== Prototypes =============================== */ /* =============================== Prototypes =============================== */
gguf_ctx *gguf_init(char *filename); gguf_ctx *gguf_init(char *filename);
int gguf_remap(gguf_ctx *ctx);
void gguf_rewind(gguf_ctx *ctx);
void gguf_end(gguf_ctx *ctx); void gguf_end(gguf_ctx *ctx);
int gguf_get_key(gguf_ctx *ctx, gguf_key *key); int gguf_get_key(gguf_ctx *ctx, gguf_key *key);
int gguf_get_tensor(gguf_ctx *ctx, gguf_tensor *tensor); int gguf_get_tensor(gguf_ctx *ctx, gguf_tensor *tensor);