1289 lines
45 KiB
C
1289 lines
45 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
||
/* Copyright (C) 2026 David (Corey) Herrera */
|
||
#include "convert.h"
|
||
#include "../safetensors/safetensors.h"
|
||
#include "../quantization/quantization.h"
|
||
#include "../tokenizer/tokenizer.h"
|
||
|
||
#include <stdio.h>
|
||
#include <stdlib.h>
|
||
#include <string.h>
|
||
#include <math.h>
|
||
|
||
/* ============================================================= */
|
||
/* Minimal JSON helpers (same pattern as safetensors.c) */
|
||
/* ============================================================= */
|
||
|
||
static char* read_file_str(const char *path) {
|
||
FILE *f = fopen(path, "r");
|
||
if (!f) return NULL;
|
||
fseek(f, 0, SEEK_END);
|
||
long size = ftell(f);
|
||
fseek(f, 0, SEEK_SET);
|
||
char *buf = malloc(size + 1);
|
||
size_t rd = fread(buf, 1, size, f);
|
||
buf[rd] = '\0';
|
||
fclose(f);
|
||
return buf;
|
||
}
|
||
|
||
/* Extract integer value for a JSON key like "num_hidden_layers": 28 */
|
||
static int json_get_int(const char *json, const char *key, int default_val) {
|
||
char search[128];
|
||
snprintf(search, sizeof(search), "\"%s\"", key);
|
||
const char *pos = strstr(json, search);
|
||
if (!pos) return default_val;
|
||
pos += strlen(search);
|
||
while (*pos && (*pos == ' ' || *pos == ':' || *pos == '\t' || *pos == '\n')) pos++;
|
||
return (int)strtol(pos, NULL, 10);
|
||
}
|
||
|
||
/* Extract float value for a JSON key */
|
||
static float json_get_float(const char *json, const char *key, float default_val) {
|
||
char search[128];
|
||
snprintf(search, sizeof(search), "\"%s\"", key);
|
||
const char *pos = strstr(json, search);
|
||
if (!pos) return default_val;
|
||
pos += strlen(search);
|
||
while (*pos && (*pos == ' ' || *pos == ':' || *pos == '\t' || *pos == '\n')) pos++;
|
||
return strtof(pos, NULL);
|
||
}
|
||
|
||
/* Extract string value for a JSON key (caller must free) */
|
||
static char* json_get_string(const char *json, const char *key) {
|
||
char search[128];
|
||
snprintf(search, sizeof(search), "\"%s\"", key);
|
||
const char *pos = strstr(json, search);
|
||
if (!pos) return NULL;
|
||
pos += strlen(search);
|
||
while (*pos && (*pos == ' ' || *pos == ':' || *pos == '\t' || *pos == '\n')) pos++;
|
||
if (*pos != '"') return NULL;
|
||
pos++;
|
||
const char *end = pos;
|
||
while (*end && *end != '"') {
|
||
if (*end == '\\') end++;
|
||
if (*end) end++;
|
||
}
|
||
size_t len = end - pos;
|
||
char *result = malloc(len + 1);
|
||
memcpy(result, pos, len);
|
||
result[len] = '\0';
|
||
return result;
|
||
}
|
||
|
||
/* Extract boolean value for a JSON key */
|
||
static int json_get_bool(const char *json, const char *key, int default_val) {
|
||
char search[128];
|
||
snprintf(search, sizeof(search), "\"%s\"", key);
|
||
const char *pos = strstr(json, search);
|
||
if (!pos) return default_val;
|
||
pos += strlen(search);
|
||
while (*pos && (*pos == ' ' || *pos == ':' || *pos == '\t' || *pos == '\n')) pos++;
|
||
if (*pos == 't') return 1;
|
||
if (*pos == 'f') return 0;
|
||
return default_val;
|
||
}
|
||
|
||
/* ============================================================= */
|
||
/* Architecture Detection */
|
||
/* ============================================================= */
|
||
|
||
const char* fc_arch_name(fc_arch arch) {
|
||
switch (arch) {
|
||
case FC_ARCH_LLAMA: return "Llama";
|
||
case FC_ARCH_QWEN2: return "Qwen2";
|
||
case FC_ARCH_QWEN3: return "Qwen3";
|
||
case FC_ARCH_PHI3: return "Phi-3";
|
||
case FC_ARCH_GEMMA2: return "Gemma2";
|
||
case FC_ARCH_UNKNOWN: return "Unknown";
|
||
}
|
||
return "Unknown";
|
||
}
|
||
|
||
static fc_arch detect_arch(const char *model_type) {
|
||
if (!model_type) return FC_ARCH_UNKNOWN;
|
||
if (strstr(model_type, "qwen3")) return FC_ARCH_QWEN3;
|
||
if (strstr(model_type, "qwen2") || strstr(model_type, "qwen")) return FC_ARCH_QWEN2;
|
||
if (strstr(model_type, "llama") || strstr(model_type, "codellama")) return FC_ARCH_LLAMA;
|
||
if (strstr(model_type, "mistral")) return FC_ARCH_LLAMA;
|
||
if (strstr(model_type, "phi3") || strstr(model_type, "phi")) return FC_ARCH_PHI3;
|
||
if (strstr(model_type, "gemma2") || strstr(model_type, "gemma")) return FC_ARCH_GEMMA2;
|
||
return FC_ARCH_UNKNOWN;
|
||
}
|
||
|
||
/* ============================================================= */
|
||
/* Config Parsing */
|
||
/* ============================================================= */
|
||
|
||
fc_model_config* fc_parse_config(const char *path) {
|
||
char *json = read_file_str(path);
|
||
if (!json) {
|
||
fprintf(stderr, "ERROR: cannot read %s\n", path);
|
||
return NULL;
|
||
}
|
||
|
||
fc_model_config *cfg = calloc(1, sizeof(fc_model_config));
|
||
|
||
/* Extract model_type */
|
||
char *mt = json_get_string(json, "model_type");
|
||
if (mt) {
|
||
snprintf(cfg->model_type, sizeof(cfg->model_type), "%s", mt);
|
||
cfg->arch = detect_arch(mt);
|
||
free(mt);
|
||
} else {
|
||
cfg->arch = FC_ARCH_UNKNOWN;
|
||
strcpy(cfg->model_type, "unknown");
|
||
}
|
||
|
||
/* Core dimensions */
|
||
cfg->n_layers = json_get_int(json, "num_hidden_layers", 28);
|
||
cfg->n_embd = json_get_int(json, "hidden_size", 1024);
|
||
cfg->n_heads = json_get_int(json, "num_attention_heads", 16);
|
||
cfg->n_kv_heads = json_get_int(json, "num_key_value_heads", cfg->n_heads);
|
||
cfg->ffn_dim = json_get_int(json, "intermediate_size", 2816);
|
||
cfg->vocab_size = json_get_int(json, "vocab_size", 32000);
|
||
cfg->max_seq_len = json_get_int(json, "max_position_embeddings", 4096);
|
||
|
||
/* Float params */
|
||
cfg->rope_theta = json_get_float(json, "rope_theta", 10000.0f);
|
||
cfg->rms_eps = json_get_float(json, "rms_norm_eps", 1e-6f);
|
||
|
||
/* Derived */
|
||
cfg->head_dim = cfg->n_embd / cfg->n_heads;
|
||
|
||
/* Optional features */
|
||
cfg->tie_word_embeddings = json_get_bool(json, "tie_word_embeddings", 1);
|
||
|
||
/* Architecture-specific defaults */
|
||
switch (cfg->arch) {
|
||
case FC_ARCH_QWEN3:
|
||
cfg->has_qk_norm = 1;
|
||
cfg->has_qkv_bias = 0;
|
||
break;
|
||
case FC_ARCH_QWEN2:
|
||
case FC_ARCH_PHI3:
|
||
cfg->has_qk_norm = 0;
|
||
cfg->has_qkv_bias = 1;
|
||
break;
|
||
case FC_ARCH_GEMMA2:
|
||
cfg->has_qk_norm = 1;
|
||
cfg->has_qkv_bias = 0;
|
||
break;
|
||
default:
|
||
cfg->has_qk_norm = 0;
|
||
cfg->has_qkv_bias = 0;
|
||
break;
|
||
}
|
||
|
||
free(json);
|
||
return cfg;
|
||
}
|
||
|
||
/* ============================================================= */
|
||
/* Data Type Conversion */
|
||
/* ============================================================= */
|
||
|
||
static void bf16_to_f32(const uint16_t *src, float *dst, int64_t n) {
|
||
for (int64_t i = 0; i < n; i++) {
|
||
uint32_t bits = (uint32_t)src[i] << 16;
|
||
memcpy(&dst[i], &bits, 4);
|
||
}
|
||
}
|
||
|
||
static void f16_to_f32(const uint16_t *src, float *dst, int64_t n) {
|
||
for (int64_t i = 0; i < n; i++) {
|
||
uint16_t h = src[i];
|
||
uint32_t sign = (uint32_t)(h >> 15) << 31;
|
||
uint32_t exp = (h >> 10) & 0x1F;
|
||
uint32_t mant = h & 0x3FF;
|
||
uint32_t bits;
|
||
if (exp == 0) {
|
||
if (mant == 0) {
|
||
bits = sign;
|
||
} else {
|
||
/* Subnormal: normalize */
|
||
exp = 1;
|
||
while (!(mant & 0x400)) { mant <<= 1; exp--; }
|
||
mant &= 0x3FF;
|
||
bits = sign | ((uint32_t)(exp + 127 - 15) << 23) | ((uint32_t)mant << 13);
|
||
}
|
||
} else if (exp == 31) {
|
||
bits = sign | 0x7F800000 | ((uint32_t)mant << 13);
|
||
} else {
|
||
bits = sign | ((uint32_t)(exp + 127 - 15) << 23) | ((uint32_t)mant << 13);
|
||
}
|
||
memcpy(&dst[i], &bits, 4);
|
||
}
|
||
}
|
||
|
||
static void bf16_to_f16(const uint16_t *src, uint16_t *dst, int64_t n) {
|
||
/* BF16 → F32 → F16 (simple, not speed-critical) */
|
||
for (int64_t i = 0; i < n; i++) {
|
||
uint32_t bits = (uint32_t)src[i] << 16;
|
||
float f;
|
||
memcpy(&f, &bits, 4);
|
||
/* f32 to f16: truncate */
|
||
uint32_t fbits;
|
||
memcpy(&fbits, &f, 4);
|
||
uint32_t sign = (fbits >> 16) & 0x8000;
|
||
int32_t exp = ((fbits >> 23) & 0xFF) - 127 + 15;
|
||
uint32_t mant = (fbits >> 13) & 0x3FF;
|
||
if (exp <= 0) {
|
||
dst[i] = (uint16_t)sign; /* flush to zero */
|
||
} else if (exp >= 31) {
|
||
dst[i] = (uint16_t)(sign | 0x7C00); /* infinity */
|
||
} else {
|
||
dst[i] = (uint16_t)(sign | (exp << 10) | mant);
|
||
}
|
||
}
|
||
}
|
||
|
||
/*
|
||
* Convert safetensors tensor data to f32.
|
||
* Handles BF16, F16, F32 input types. Returns malloc'd f32 array.
|
||
* Caller must free.
|
||
*/
|
||
static float* tensor_to_f32(const void *data, fs_dtype dtype, int64_t numel) {
|
||
float *f32 = malloc(numel * sizeof(float));
|
||
if (!f32) return NULL;
|
||
|
||
switch (dtype) {
|
||
case FS_DTYPE_BF16:
|
||
bf16_to_f32((const uint16_t *)data, f32, numel);
|
||
break;
|
||
case FS_DTYPE_F16:
|
||
f16_to_f32((const uint16_t *)data, f32, numel);
|
||
break;
|
||
case FS_DTYPE_F32:
|
||
memcpy(f32, data, numel * sizeof(float));
|
||
break;
|
||
default:
|
||
fprintf(stderr, "WARNING: unsupported dtype %d, zeroing\n", dtype);
|
||
memset(f32, 0, numel * sizeof(float));
|
||
break;
|
||
}
|
||
return f32;
|
||
}
|
||
|
||
/*
|
||
* Convert safetensors tensor data to f16 (for embeddings).
|
||
* Handles BF16, F16, F32 input. Returns malloc'd uint16_t array.
|
||
* Caller must free.
|
||
*/
|
||
static uint16_t* tensor_to_f16(const void *data, fs_dtype dtype, int64_t numel) {
|
||
uint16_t *f16 = malloc(numel * sizeof(uint16_t));
|
||
if (!f16) return NULL;
|
||
|
||
switch (dtype) {
|
||
case FS_DTYPE_BF16:
|
||
bf16_to_f16((const uint16_t *)data, f16, numel);
|
||
break;
|
||
case FS_DTYPE_F16:
|
||
memcpy(f16, data, numel * sizeof(uint16_t));
|
||
break;
|
||
case FS_DTYPE_F32: {
|
||
const float *f32 = (const float *)data;
|
||
for (int64_t i = 0; i < numel; i++) {
|
||
uint32_t fbits;
|
||
memcpy(&fbits, &f32[i], 4);
|
||
uint32_t sign = (fbits >> 16) & 0x8000;
|
||
int32_t exp = ((fbits >> 23) & 0xFF) - 127 + 15;
|
||
uint32_t mant = (fbits >> 13) & 0x3FF;
|
||
if (exp <= 0) f16[i] = (uint16_t)sign;
|
||
else if (exp >= 31) f16[i] = (uint16_t)(sign | 0x7C00);
|
||
else f16[i] = (uint16_t)(sign | (exp << 10) | mant);
|
||
}
|
||
break;
|
||
}
|
||
default:
|
||
memset(f16, 0, numel * sizeof(uint16_t));
|
||
break;
|
||
}
|
||
return f16;
|
||
}
|
||
|
||
/* ============================================================= */
|
||
/* .btw Header (must match bootai/model/loader.c) */
|
||
/* ============================================================= */
|
||
|
||
#pragma pack(push, 1)
|
||
typedef struct {
|
||
char magic[4]; /* "BTAI" */
|
||
uint32_t version; /* 1 */
|
||
uint32_t arch; /* 0=qwen3, 1=qwen2, 2=llama, 3=rwkvx */
|
||
uint32_t quant; /* 0=f32, 1=f16, 2=q8, 3=q4 */
|
||
uint32_t n_layers;
|
||
uint32_t n_embd;
|
||
uint32_t n_heads;
|
||
uint32_t n_kv_heads;
|
||
uint32_t head_dim;
|
||
uint32_t ffn_dim;
|
||
uint32_t vocab_size;
|
||
float rope_theta;
|
||
float rms_eps;
|
||
uint32_t max_seq_len;
|
||
uint32_t has_qk_norm;
|
||
uint32_t has_qkv_bias;
|
||
uint8_t reserved[64];
|
||
} btw_header;
|
||
#pragma pack(pop)
|
||
|
||
static uint32_t arch_to_btw(fc_arch arch) {
|
||
switch (arch) {
|
||
case FC_ARCH_QWEN3: return 0;
|
||
case FC_ARCH_QWEN2: return 1;
|
||
case FC_ARCH_LLAMA: return 2;
|
||
case FC_ARCH_PHI3: return 1; /* same layout as Qwen2 (has bias) */
|
||
case FC_ARCH_GEMMA2: return 0; /* same layout as Qwen3 (has qk-norm) */
|
||
default: return 2; /* default to Llama */
|
||
}
|
||
}
|
||
|
||
/* ============================================================= */
|
||
/* Weight Export */
|
||
/* ============================================================= */
|
||
|
||
/*
|
||
* Fast per-row symmetric INT8 quantization for 2D matrices.
|
||
* Much faster than fq_quantize_int8_per_channel which does O(channels × numel).
|
||
* This does O(numel) — a single pass per row.
|
||
*/
|
||
static void quantize_per_row_q8(const float *f32, int8_t *q8, float *scales,
|
||
int rows, int cols) {
|
||
for (int r = 0; r < rows; r++) {
|
||
const float *row = f32 + (int64_t)r * cols;
|
||
float abs_max = 0.0f;
|
||
for (int c = 0; c < cols; c++) {
|
||
float a = fabsf(row[c]);
|
||
if (a > abs_max) abs_max = a;
|
||
}
|
||
float scale = abs_max / 127.0f;
|
||
scales[r] = scale;
|
||
float inv_scale = (abs_max > 0.0f) ? 127.0f / abs_max : 0.0f;
|
||
int8_t *qrow = q8 + (int64_t)r * cols;
|
||
for (int c = 0; c < cols; c++) {
|
||
int v = (int)roundf(row[c] * inv_scale);
|
||
if (v > 127) v = 127;
|
||
if (v < -127) v = -127;
|
||
qrow[c] = (int8_t)v;
|
||
}
|
||
}
|
||
}
|
||
|
||
/* Write a Q8 quantized weight matrix to file.
|
||
* Reads raw data from safetensors, converts to f32, quantizes per-row, writes. */
|
||
static bool write_q8_weight(FILE *out, fs_file *sf, const char *name,
|
||
int rows, int cols, uint64_t *total) {
|
||
fs_tensor_info *info = fs_get_info(sf, name);
|
||
if (!info) {
|
||
fprintf(stderr, " WARNING: tensor '%s' not found, writing zeros\n", name);
|
||
int64_t data_size = (int64_t)rows * cols;
|
||
int8_t *zeros = calloc(data_size, 1);
|
||
float *zero_scales = calloc(rows, sizeof(float));
|
||
fwrite(zeros, 1, data_size, out);
|
||
fwrite(zero_scales, sizeof(float), rows, out);
|
||
free(zeros);
|
||
free(zero_scales);
|
||
*total += data_size + rows * sizeof(float);
|
||
return true;
|
||
}
|
||
|
||
const void *raw = fs_get_data(sf, name);
|
||
int64_t numel = (int64_t)rows * cols;
|
||
float *f32 = tensor_to_f32(raw, info->dtype, numel);
|
||
fs_free_info(info);
|
||
|
||
if (!f32) {
|
||
fprintf(stderr, " ERROR: failed to convert '%s' to f32\n", name);
|
||
return false;
|
||
}
|
||
|
||
/* Fast per-row quantization */
|
||
int8_t *q8 = malloc(numel);
|
||
float *scales = malloc(rows * sizeof(float));
|
||
quantize_per_row_q8(f32, q8, scales, rows, cols);
|
||
free(f32);
|
||
|
||
fwrite(q8, 1, numel, out);
|
||
fwrite(scales, sizeof(float), rows, out);
|
||
*total += numel + rows * sizeof(float);
|
||
|
||
free(q8);
|
||
free(scales);
|
||
return true;
|
||
}
|
||
|
||
/* Write an f32 weight to file. Converts from BF16/F16 if needed. */
|
||
static bool write_f32_weight(FILE *out, fs_file *sf, const char *name,
|
||
int64_t numel, uint64_t *total) {
|
||
fs_tensor_info *info = fs_get_info(sf, name);
|
||
if (!info) {
|
||
fprintf(stderr, " WARNING: tensor '%s' not found, writing zeros\n", name);
|
||
float *zeros = calloc(numel, sizeof(float));
|
||
fwrite(zeros, sizeof(float), numel, out);
|
||
free(zeros);
|
||
*total += numel * sizeof(float);
|
||
return true;
|
||
}
|
||
|
||
const void *raw = fs_get_data(sf, name);
|
||
float *f32 = tensor_to_f32(raw, info->dtype, numel);
|
||
fs_free_info(info);
|
||
|
||
fwrite(f32, sizeof(float), numel, out);
|
||
*total += numel * sizeof(float);
|
||
free(f32);
|
||
return true;
|
||
}
|
||
|
||
/* Write an f16 embedding/lm_head weight. Converts from BF16 if needed. */
|
||
static bool write_f16_weight(FILE *out, fs_file *sf, const char *name,
|
||
int64_t numel, uint64_t *total) {
|
||
fs_tensor_info *info = fs_get_info(sf, name);
|
||
if (!info) {
|
||
fprintf(stderr, " WARNING: tensor '%s' not found\n", name);
|
||
return false;
|
||
}
|
||
|
||
const void *raw = fs_get_data(sf, name);
|
||
uint16_t *f16 = tensor_to_f16(raw, info->dtype, numel);
|
||
fs_free_info(info);
|
||
|
||
fwrite(f16, sizeof(uint16_t), numel, out);
|
||
*total += numel * sizeof(uint16_t);
|
||
free(f16);
|
||
return true;
|
||
}
|
||
|
||
bool fc_export_weights(const char *safetensors_path,
|
||
fc_model_config *config,
|
||
const char *output_path) {
|
||
fprintf(stderr, "Opening: %s\n", safetensors_path);
|
||
fs_file *sf = fs_open(safetensors_path);
|
||
if (!sf) {
|
||
fprintf(stderr, "ERROR: cannot open %s\n", safetensors_path);
|
||
return false;
|
||
}
|
||
|
||
int nt = fs_num_tensors(sf);
|
||
fprintf(stderr, " %d tensors found\n", nt);
|
||
|
||
/* Detect naming prefix: "model.layers..." vs "layers..." */
|
||
char prefix[64] = "";
|
||
if (fs_has_tensor(sf, "model.embed_tokens.weight")) {
|
||
strcpy(prefix, "model.");
|
||
}
|
||
|
||
FILE *out = fopen(output_path, "wb");
|
||
if (!out) {
|
||
fprintf(stderr, "ERROR: cannot create %s\n", output_path);
|
||
fs_close(sf);
|
||
return false;
|
||
}
|
||
|
||
/* Write header */
|
||
btw_header hdr;
|
||
memset(&hdr, 0, sizeof(hdr));
|
||
memcpy(hdr.magic, "BTAI", 4);
|
||
hdr.version = 1;
|
||
hdr.arch = arch_to_btw(config->arch);
|
||
hdr.quant = 1; /* F16 */
|
||
hdr.n_layers = config->n_layers;
|
||
hdr.n_embd = config->n_embd;
|
||
hdr.n_heads = config->n_heads;
|
||
hdr.n_kv_heads = config->n_kv_heads;
|
||
hdr.head_dim = config->head_dim;
|
||
hdr.ffn_dim = config->ffn_dim;
|
||
hdr.vocab_size = config->vocab_size;
|
||
hdr.rope_theta = config->rope_theta;
|
||
hdr.rms_eps = config->rms_eps;
|
||
hdr.max_seq_len = config->max_seq_len;
|
||
hdr.has_qk_norm = config->has_qk_norm;
|
||
hdr.has_qkv_bias = config->has_qkv_bias;
|
||
|
||
fwrite(&hdr, sizeof(hdr), 1, out);
|
||
uint64_t total = sizeof(hdr);
|
||
|
||
fprintf(stderr, "\nModel: %s (%s)\n", config->model_type, fc_arch_name(config->arch));
|
||
fprintf(stderr, " Layers: %d, Hidden: %d, Heads: %d/%d, FFN: %d\n",
|
||
config->n_layers, config->n_embd, config->n_heads,
|
||
config->n_kv_heads, config->ffn_dim);
|
||
fprintf(stderr, " QK-Norm: %d, QKV-Bias: %d, Tied: %d\n",
|
||
config->has_qk_norm, config->has_qkv_bias,
|
||
config->tie_word_embeddings);
|
||
|
||
int q_dim = config->n_heads * config->head_dim;
|
||
int kv_dim = config->n_kv_heads * config->head_dim;
|
||
|
||
/* 1. Embedding: [vocab, n_embd] as f16 */
|
||
char tname[256];
|
||
snprintf(tname, sizeof(tname), "%sembed_tokens.weight", prefix);
|
||
int64_t embed_numel = (int64_t)config->vocab_size * config->n_embd;
|
||
fprintf(stderr, "\nWriting embed_tokens [%d x %d] f16...\n",
|
||
config->vocab_size, config->n_embd);
|
||
write_f16_weight(out, sf, tname, embed_numel, &total);
|
||
|
||
/* 2. Per-layer weights */
|
||
for (int i = 0; i < config->n_layers; i++) {
|
||
char lp[128];
|
||
snprintf(lp, sizeof(lp), "%slayers.%d.", prefix, i);
|
||
|
||
/* input_layernorm.weight: [n_embd] f32 */
|
||
snprintf(tname, sizeof(tname), "%sinput_layernorm.weight", lp);
|
||
write_f32_weight(out, sf, tname, config->n_embd, &total);
|
||
|
||
/* Q projection: [q_dim, n_embd] F16 */
|
||
snprintf(tname, sizeof(tname), "%sself_attn.q_proj.weight", lp);
|
||
write_f16_weight(out, sf, tname, (int64_t)q_dim * config->n_embd, &total);
|
||
|
||
/* K projection: [kv_dim, n_embd] F16 */
|
||
snprintf(tname, sizeof(tname), "%sself_attn.k_proj.weight", lp);
|
||
write_f16_weight(out, sf, tname, (int64_t)kv_dim * config->n_embd, &total);
|
||
|
||
/* V projection: [kv_dim, n_embd] F16 */
|
||
snprintf(tname, sizeof(tname), "%sself_attn.v_proj.weight", lp);
|
||
write_f16_weight(out, sf, tname, (int64_t)kv_dim * config->n_embd, &total);
|
||
|
||
/* O projection: [n_embd, q_dim] F16 */
|
||
snprintf(tname, sizeof(tname), "%sself_attn.o_proj.weight", lp);
|
||
write_f16_weight(out, sf, tname, (int64_t)config->n_embd * q_dim, &total);
|
||
|
||
/* QK-Norm (if present) */
|
||
if (config->has_qk_norm) {
|
||
snprintf(tname, sizeof(tname), "%sself_attn.q_norm.weight", lp);
|
||
write_f32_weight(out, sf, tname, config->head_dim, &total);
|
||
snprintf(tname, sizeof(tname), "%sself_attn.k_norm.weight", lp);
|
||
write_f32_weight(out, sf, tname, config->head_dim, &total);
|
||
}
|
||
|
||
/* QKV biases (if present) */
|
||
if (config->has_qkv_bias) {
|
||
snprintf(tname, sizeof(tname), "%sself_attn.q_proj.bias", lp);
|
||
write_f32_weight(out, sf, tname, q_dim, &total);
|
||
snprintf(tname, sizeof(tname), "%sself_attn.k_proj.bias", lp);
|
||
write_f32_weight(out, sf, tname, kv_dim, &total);
|
||
snprintf(tname, sizeof(tname), "%sself_attn.v_proj.bias", lp);
|
||
write_f32_weight(out, sf, tname, kv_dim, &total);
|
||
}
|
||
|
||
/* post_attention_layernorm.weight: [n_embd] f32 */
|
||
snprintf(tname, sizeof(tname), "%spost_attention_layernorm.weight", lp);
|
||
write_f32_weight(out, sf, tname, config->n_embd, &total);
|
||
|
||
/* gate_proj: [ffn_dim, n_embd] F16 */
|
||
snprintf(tname, sizeof(tname), "%smlp.gate_proj.weight", lp);
|
||
write_f16_weight(out, sf, tname, (int64_t)config->ffn_dim * config->n_embd, &total);
|
||
|
||
/* up_proj: [ffn_dim, n_embd] F16 */
|
||
snprintf(tname, sizeof(tname), "%smlp.up_proj.weight", lp);
|
||
write_f16_weight(out, sf, tname, (int64_t)config->ffn_dim * config->n_embd, &total);
|
||
|
||
/* down_proj: [n_embd, ffn_dim] F16 */
|
||
snprintf(tname, sizeof(tname), "%smlp.down_proj.weight", lp);
|
||
write_f16_weight(out, sf, tname, (int64_t)config->n_embd * config->ffn_dim, &total);
|
||
|
||
if ((i + 1) % 4 == 0 || i == config->n_layers - 1) {
|
||
fprintf(stderr, " Layer %d/%d done (%llu MB)\n",
|
||
i + 1, config->n_layers, (unsigned long long)(total / 1024 / 1024));
|
||
}
|
||
}
|
||
|
||
/* 3. Final norm: [n_embd] f32 */
|
||
snprintf(tname, sizeof(tname), "%snorm.weight", prefix);
|
||
write_f32_weight(out, sf, tname, config->n_embd, &total);
|
||
|
||
/* 4. LM head: [vocab, n_embd] f16 (skip if tied with embed_tokens) */
|
||
if (!config->tie_word_embeddings) {
|
||
fprintf(stderr, "\nWriting lm_head [%d x %d] f16...\n",
|
||
config->vocab_size, config->n_embd);
|
||
write_f16_weight(out, sf, "lm_head.weight", embed_numel, &total);
|
||
} else {
|
||
fprintf(stderr, "\n lm_head: tied with embed_tokens (shared)\n");
|
||
}
|
||
|
||
fclose(out);
|
||
fs_close(sf);
|
||
|
||
fprintf(stderr, "\nDone! Total: %.1f MB -> %s\n",
|
||
(double)total / 1024.0 / 1024.0, output_path);
|
||
return true;
|
||
}
|
||
|
||
/* ============================================================= */
|
||
/* .btv Vocab Export */
|
||
/* ============================================================= */
|
||
|
||
#pragma pack(push, 1)
|
||
typedef struct {
|
||
char magic[4]; /* "BTAV" */
|
||
uint32_t n_tokens;
|
||
uint32_t n_merges;
|
||
uint32_t n_special;
|
||
uint8_t reserved[16];
|
||
} btv_header;
|
||
#pragma pack(pop)
|
||
|
||
/*
|
||
* Fast direct vocab.json parser.
|
||
* Format: {"token_str": id, "another": id2, ...}
|
||
* Handles JSON escapes: \n \t \r \\ \" \uXXXX
|
||
* Returns malloc'd array of strings indexed by token ID.
|
||
* Caller must free each string and the array.
|
||
*/
|
||
static char** parse_vocab_json(const char *path, int *n_tokens_out) {
|
||
char *json = read_file_str(path);
|
||
if (!json) return NULL;
|
||
|
||
/* First pass: find max token ID to size the array */
|
||
int max_id = -1;
|
||
int count = 0;
|
||
const char *p = json;
|
||
while (*p) {
|
||
/* Find next quoted key */
|
||
while (*p && *p != '"') p++;
|
||
if (!*p) break;
|
||
p++; /* skip opening quote */
|
||
/* Skip the key string */
|
||
while (*p && *p != '"') {
|
||
if (*p == '\\') p++;
|
||
if (*p) p++;
|
||
}
|
||
if (!*p) break;
|
||
p++; /* skip closing quote */
|
||
/* Find colon and number */
|
||
while (*p && *p != ':') p++;
|
||
if (!*p) break;
|
||
p++; /* skip colon */
|
||
while (*p == ' ' || *p == '\t' || *p == '\n' || *p == '\r') p++;
|
||
int id = (int)strtol(p, NULL, 10);
|
||
if (id > max_id) max_id = id;
|
||
count++;
|
||
/* Skip past the number */
|
||
while (*p && *p != ',' && *p != '}') p++;
|
||
}
|
||
|
||
if (max_id < 0) {
|
||
free(json);
|
||
*n_tokens_out = 0;
|
||
return NULL;
|
||
}
|
||
|
||
int n_tokens = max_id + 1;
|
||
char **tokens = calloc(n_tokens, sizeof(char *));
|
||
|
||
/* Second pass: extract key strings and assign to their IDs */
|
||
p = json;
|
||
while (*p) {
|
||
while (*p && *p != '"') p++;
|
||
if (!*p) break;
|
||
p++; /* skip opening quote */
|
||
|
||
/* Read key with escape handling */
|
||
char key_buf[1024];
|
||
int ki = 0;
|
||
while (*p && *p != '"' && ki < (int)sizeof(key_buf) - 1) {
|
||
if (*p == '\\') {
|
||
p++;
|
||
switch (*p) {
|
||
case 'n': key_buf[ki++] = '\n'; break;
|
||
case 't': key_buf[ki++] = '\t'; break;
|
||
case 'r': key_buf[ki++] = '\r'; break;
|
||
case '\\': key_buf[ki++] = '\\'; break;
|
||
case '"': key_buf[ki++] = '"'; break;
|
||
case '/': key_buf[ki++] = '/'; break;
|
||
case 'u': {
|
||
/* \uXXXX -> UTF-8 */
|
||
if (p[1] && p[2] && p[3] && p[4]) {
|
||
char hex[5] = {p[1], p[2], p[3], p[4], 0};
|
||
uint32_t cp = (uint32_t)strtoul(hex, NULL, 16);
|
||
p += 4;
|
||
if (cp < 0x80) {
|
||
key_buf[ki++] = (char)cp;
|
||
} else if (cp < 0x800) {
|
||
key_buf[ki++] = (char)(0xC0 | (cp >> 6));
|
||
key_buf[ki++] = (char)(0x80 | (cp & 0x3F));
|
||
} else {
|
||
key_buf[ki++] = (char)(0xE0 | (cp >> 12));
|
||
key_buf[ki++] = (char)(0x80 | ((cp >> 6) & 0x3F));
|
||
key_buf[ki++] = (char)(0x80 | (cp & 0x3F));
|
||
}
|
||
}
|
||
break;
|
||
}
|
||
default: key_buf[ki++] = *p; break;
|
||
}
|
||
if (*p) p++;
|
||
} else {
|
||
key_buf[ki++] = *p++;
|
||
}
|
||
}
|
||
key_buf[ki] = '\0';
|
||
if (!*p) break;
|
||
p++; /* skip closing quote */
|
||
|
||
/* Find colon and number */
|
||
while (*p && *p != ':') p++;
|
||
if (!*p) break;
|
||
p++;
|
||
while (*p == ' ' || *p == '\t' || *p == '\n' || *p == '\r') p++;
|
||
int id = (int)strtol(p, NULL, 10);
|
||
|
||
if (id >= 0 && id < n_tokens) {
|
||
tokens[id] = strdup(key_buf);
|
||
}
|
||
|
||
while (*p && *p != ',' && *p != '}') p++;
|
||
}
|
||
|
||
free(json);
|
||
*n_tokens_out = n_tokens;
|
||
return tokens;
|
||
}
|
||
|
||
/*
|
||
* Fast merges.txt parser.
|
||
* Format: optional "#version: 0.2\n" header, then "token_a token_b\n" per line.
|
||
* Returns two malloc'd arrays of strings (a[] and b[]) and count.
|
||
*/
|
||
static void parse_merges_txt(const char *path, char ***a_out, char ***b_out,
|
||
int *n_merges_out) {
|
||
char *data = read_file_str(path);
|
||
if (!data) {
|
||
*a_out = NULL;
|
||
*b_out = NULL;
|
||
*n_merges_out = 0;
|
||
return;
|
||
}
|
||
|
||
/* Count lines (excluding header and empty lines) */
|
||
int n_lines = 0;
|
||
const char *p = data;
|
||
while (*p) {
|
||
while (*p == '\n' || *p == '\r') p++;
|
||
if (!*p) break;
|
||
if (*p != '#') n_lines++;
|
||
while (*p && *p != '\n') p++;
|
||
}
|
||
|
||
char **a = malloc(n_lines * sizeof(char *));
|
||
char **b = malloc(n_lines * sizeof(char *));
|
||
int count = 0;
|
||
|
||
p = data;
|
||
while (*p) {
|
||
/* Skip whitespace/newlines */
|
||
while (*p == '\n' || *p == '\r') p++;
|
||
if (!*p) break;
|
||
|
||
/* Skip comment lines */
|
||
if (*p == '#') {
|
||
while (*p && *p != '\n') p++;
|
||
continue;
|
||
}
|
||
|
||
/* Read token_a (up to space) */
|
||
const char *start = p;
|
||
while (*p && *p != ' ' && *p != '\n' && *p != '\r') p++;
|
||
size_t a_len = p - start;
|
||
|
||
if (*p != ' ') {
|
||
/* Malformed line, skip */
|
||
while (*p && *p != '\n') p++;
|
||
continue;
|
||
}
|
||
p++; /* skip space */
|
||
|
||
/* Read token_b (up to newline) */
|
||
const char *b_start = p;
|
||
while (*p && *p != '\n' && *p != '\r') p++;
|
||
size_t b_len = p - b_start;
|
||
|
||
/* Trim trailing spaces from b */
|
||
while (b_len > 0 && (b_start[b_len - 1] == ' ' || b_start[b_len - 1] == '\r'))
|
||
b_len--;
|
||
|
||
a[count] = strndup(start, a_len);
|
||
b[count] = strndup(b_start, b_len);
|
||
count++;
|
||
}
|
||
|
||
free(data);
|
||
*a_out = a;
|
||
*b_out = b;
|
||
*n_merges_out = count;
|
||
}
|
||
|
||
/*
|
||
* Simple hash map: string → int.
|
||
* Uses FNV-1a hash with open addressing (linear probing).
|
||
* For 151K entries, 256K buckets gives ~60% load factor.
|
||
*/
|
||
#define VOCAB_MAP_CAPACITY (1 << 18) /* 262144 buckets */
|
||
#define VOCAB_MAP_MASK (VOCAB_MAP_CAPACITY - 1)
|
||
|
||
typedef struct {
|
||
const char *key;
|
||
int value;
|
||
} vocab_map_entry;
|
||
|
||
typedef struct {
|
||
vocab_map_entry *buckets;
|
||
} vocab_map;
|
||
|
||
static uint32_t fnv1a(const char *str) {
|
||
uint32_t h = 2166136261u;
|
||
while (*str) {
|
||
h ^= (uint8_t)*str++;
|
||
h *= 16777619u;
|
||
}
|
||
return h;
|
||
}
|
||
|
||
static vocab_map* vocab_map_create(void) {
|
||
vocab_map *m = calloc(1, sizeof(vocab_map));
|
||
m->buckets = calloc(VOCAB_MAP_CAPACITY, sizeof(vocab_map_entry));
|
||
return m;
|
||
}
|
||
|
||
static void vocab_map_put(vocab_map *m, const char *key, int value) {
|
||
uint32_t idx = fnv1a(key) & VOCAB_MAP_MASK;
|
||
while (m->buckets[idx].key) {
|
||
if (strcmp(m->buckets[idx].key, key) == 0) {
|
||
m->buckets[idx].value = value;
|
||
return;
|
||
}
|
||
idx = (idx + 1) & VOCAB_MAP_MASK;
|
||
}
|
||
m->buckets[idx].key = key;
|
||
m->buckets[idx].value = value;
|
||
}
|
||
|
||
static int vocab_map_get(vocab_map *m, const char *key) {
|
||
uint32_t idx = fnv1a(key) & VOCAB_MAP_MASK;
|
||
while (m->buckets[idx].key) {
|
||
if (strcmp(m->buckets[idx].key, key) == 0)
|
||
return m->buckets[idx].value;
|
||
idx = (idx + 1) & VOCAB_MAP_MASK;
|
||
}
|
||
return -1;
|
||
}
|
||
|
||
static void vocab_map_free(vocab_map *m) {
|
||
if (m) { free(m->buckets); free(m); }
|
||
}
|
||
|
||
/*
|
||
* GPT-2 bytes_to_unicode() mapping.
|
||
* Maps each byte 0-255 to a unicode codepoint.
|
||
* Printable ASCII and Latin-1 supplement map to themselves.
|
||
* Control bytes and gaps map to U+0100+.
|
||
*/
|
||
static void build_gpt2_byte_map(uint32_t *byte_to_unicode, int *unicode_to_byte) {
|
||
/* Initialize reverse map to -1 */
|
||
for (int i = 0; i < 65536; i++) unicode_to_byte[i] = -1;
|
||
|
||
/* Ranges that map to themselves: !-~ (0x21-0x7E), ¡-¬ (0xA1-0xAC), ®-ÿ (0xAE-0xFF) */
|
||
int n = 0;
|
||
for (int b = 0; b < 256; b++) {
|
||
if ((b >= 0x21 && b <= 0x7E) ||
|
||
(b >= 0xA1 && b <= 0xAC) ||
|
||
(b >= 0xAE && b <= 0xFF)) {
|
||
byte_to_unicode[b] = (uint32_t)b;
|
||
unicode_to_byte[b] = b;
|
||
} else {
|
||
byte_to_unicode[b] = 256 + n;
|
||
unicode_to_byte[256 + n] = b;
|
||
n++;
|
||
}
|
||
}
|
||
}
|
||
|
||
/*
|
||
* Convert a GPT-2 unicode token string to raw bytes.
|
||
* E.g., "Ġhello" (U+0120 + hello) → " hello" (space + hello)
|
||
* Returns malloc'd string. Caller must free.
|
||
*/
|
||
static char* gpt2_unicode_to_bytes(const char *utf8_str, const int *unicode_to_byte, int *out_len) {
|
||
if (!utf8_str) { *out_len = 0; return NULL; }
|
||
|
||
/* Decode UTF-8 to codepoints, map each through unicode_to_byte */
|
||
size_t slen = strlen(utf8_str);
|
||
char *result = malloc(slen + 1); /* raw bytes can't be longer than UTF-8 */
|
||
int ri = 0;
|
||
|
||
const uint8_t *p = (const uint8_t *)utf8_str;
|
||
while (*p) {
|
||
uint32_t cp;
|
||
if (*p < 0x80) {
|
||
cp = *p++;
|
||
} else if ((*p & 0xE0) == 0xC0) {
|
||
cp = (*p & 0x1F) << 6;
|
||
p++;
|
||
if (*p) { cp |= (*p & 0x3F); p++; }
|
||
} else if ((*p & 0xF0) == 0xE0) {
|
||
cp = (*p & 0x0F) << 12;
|
||
p++;
|
||
if (*p) { cp |= (*p & 0x3F) << 6; p++; }
|
||
if (*p) { cp |= (*p & 0x3F); p++; }
|
||
} else if ((*p & 0xF8) == 0xF0) {
|
||
cp = (*p & 0x07) << 18;
|
||
p++;
|
||
if (*p) { cp |= (*p & 0x3F) << 12; p++; }
|
||
if (*p) { cp |= (*p & 0x3F) << 6; p++; }
|
||
if (*p) { cp |= (*p & 0x3F); p++; }
|
||
} else {
|
||
p++; continue; /* invalid UTF-8, skip */
|
||
}
|
||
|
||
/* Map unicode codepoint to raw byte */
|
||
if (cp < 65536 && unicode_to_byte[cp] >= 0) {
|
||
result[ri++] = (char)(uint8_t)unicode_to_byte[cp];
|
||
} else {
|
||
/* Not a GPT-2 byte character — keep as UTF-8 (for multi-codepoint tokens) */
|
||
/* Re-encode codepoint to UTF-8 */
|
||
if (cp < 0x80) {
|
||
result[ri++] = (char)cp;
|
||
} else if (cp < 0x800) {
|
||
result[ri++] = (char)(0xC0 | (cp >> 6));
|
||
result[ri++] = (char)(0x80 | (cp & 0x3F));
|
||
} else if (cp < 0x10000) {
|
||
result[ri++] = (char)(0xE0 | (cp >> 12));
|
||
result[ri++] = (char)(0x80 | ((cp >> 6) & 0x3F));
|
||
result[ri++] = (char)(0x80 | (cp & 0x3F));
|
||
}
|
||
}
|
||
}
|
||
*out_len = ri;
|
||
return result;
|
||
}
|
||
|
||
bool fc_export_vocab(const char *tokenizer_path,
|
||
const char *merges_path,
|
||
const char *vocab_path,
|
||
const char *output_path) {
|
||
/*
|
||
* Fast path: parse vocab.json + merges.txt directly.
|
||
* Bypasses ft_load_bpe() which is too slow for 151K+ token vocabs.
|
||
*/
|
||
if (vocab_path && merges_path) {
|
||
fprintf(stderr, "Loading vocab: %s + %s\n", vocab_path, merges_path);
|
||
|
||
/* Parse vocab.json directly */
|
||
int n_tokens = 0;
|
||
char **tokens = parse_vocab_json(vocab_path, &n_tokens);
|
||
if (!tokens || n_tokens == 0) {
|
||
fprintf(stderr, "ERROR: failed to parse vocab.json\n");
|
||
return false;
|
||
}
|
||
|
||
/* Parse merges.txt directly */
|
||
char **merge_a_str = NULL, **merge_b_str = NULL;
|
||
int n_merges = 0;
|
||
parse_merges_txt(merges_path, &merge_a_str, &merge_b_str, &n_merges);
|
||
|
||
fprintf(stderr, " Vocab: %d tokens, %d merges\n", n_tokens, n_merges);
|
||
|
||
/* Build GPT-2 byte mapping tables */
|
||
uint32_t byte_to_unicode[256];
|
||
int *unicode_to_byte = calloc(65536, sizeof(int));
|
||
build_gpt2_byte_map(byte_to_unicode, unicode_to_byte);
|
||
|
||
/* Convert all token strings from GPT-2 unicode to raw bytes.
|
||
* This makes the string blob contain actual bytes (e.g., space instead of Ġ)
|
||
* so decode works without reverse mapping. */
|
||
char **raw_tokens = malloc(n_tokens * sizeof(char *));
|
||
int *raw_lengths = malloc(n_tokens * sizeof(int));
|
||
for (int i = 0; i < n_tokens; i++) {
|
||
if (tokens[i]) {
|
||
raw_tokens[i] = gpt2_unicode_to_bytes(tokens[i], unicode_to_byte, &raw_lengths[i]);
|
||
} else {
|
||
raw_tokens[i] = NULL;
|
||
raw_lengths[i] = 0;
|
||
}
|
||
}
|
||
|
||
/* Build string blob from raw (decoded) token bytes */
|
||
size_t blob_size = 0;
|
||
for (int i = 0; i < n_tokens; i++) blob_size += raw_lengths[i];
|
||
|
||
uint32_t *offsets = malloc(n_tokens * sizeof(uint32_t));
|
||
uint16_t *lengths = malloc(n_tokens * sizeof(uint16_t));
|
||
char *blob = malloc(blob_size + 1);
|
||
size_t blob_pos = 0;
|
||
|
||
for (int i = 0; i < n_tokens; i++) {
|
||
offsets[i] = (uint32_t)blob_pos;
|
||
lengths[i] = (uint16_t)raw_lengths[i];
|
||
if (raw_tokens[i] && raw_lengths[i] > 0) {
|
||
memcpy(blob + blob_pos, raw_tokens[i], raw_lengths[i]);
|
||
blob_pos += raw_lengths[i];
|
||
}
|
||
}
|
||
|
||
/* Build hash map: ORIGINAL token string → ID for merge lookup
|
||
* (merges.txt uses GPT-2 unicode strings, not raw bytes) */
|
||
vocab_map *vmap = vocab_map_create();
|
||
for (int i = 0; i < n_tokens; i++) {
|
||
if (tokens[i]) vocab_map_put(vmap, tokens[i], i);
|
||
}
|
||
|
||
/* Resolve merge token IDs via hash map */
|
||
uint32_t *merge_a_ids = NULL;
|
||
uint32_t *merge_b_ids = NULL;
|
||
if (n_merges > 0) {
|
||
merge_a_ids = malloc(n_merges * sizeof(uint32_t));
|
||
merge_b_ids = malloc(n_merges * sizeof(uint32_t));
|
||
for (int i = 0; i < n_merges; i++) {
|
||
int a_id = merge_a_str[i] ? vocab_map_get(vmap, merge_a_str[i]) : -1;
|
||
int b_id = merge_b_str[i] ? vocab_map_get(vmap, merge_b_str[i]) : -1;
|
||
merge_a_ids[i] = (a_id >= 0) ? (uint32_t)a_id : 0;
|
||
merge_b_ids[i] = (b_id >= 0) ? (uint32_t)b_id : 0;
|
||
}
|
||
}
|
||
|
||
/* Collect special tokens */
|
||
typedef struct { uint32_t id; const char *name; } special_entry;
|
||
special_entry specials[8];
|
||
int n_special = 0;
|
||
|
||
const char *special_names[] = {
|
||
"<|im_start|>", "<|im_end|>", "<|endoftext|>",
|
||
"<s>", "</s>", "<|begin_of_text|>", "<|end_of_text|>",
|
||
NULL
|
||
};
|
||
for (int i = 0; special_names[i] && n_special < 8; i++) {
|
||
int id = vocab_map_get(vmap, special_names[i]);
|
||
if (id >= 0) {
|
||
specials[n_special].id = (uint32_t)id;
|
||
specials[n_special].name = special_names[i];
|
||
n_special++;
|
||
}
|
||
}
|
||
|
||
fprintf(stderr, " Special tokens: %d\n", n_special);
|
||
|
||
/* Write .btv file */
|
||
FILE *out = fopen(output_path, "wb");
|
||
if (!out) {
|
||
fprintf(stderr, "ERROR: cannot create %s\n", output_path);
|
||
goto cleanup_direct;
|
||
}
|
||
|
||
btv_header hdr;
|
||
memset(&hdr, 0, sizeof(hdr));
|
||
memcpy(hdr.magic, "BTAV", 4);
|
||
hdr.n_tokens = n_tokens;
|
||
hdr.n_merges = n_merges;
|
||
hdr.n_special = n_special;
|
||
|
||
fwrite(&hdr, sizeof(hdr), 1, out);
|
||
fwrite(offsets, sizeof(uint32_t), n_tokens, out);
|
||
fwrite(lengths, sizeof(uint16_t), n_tokens, out);
|
||
fwrite(blob, 1, blob_pos, out);
|
||
|
||
/* Pre-compute merge result IDs: for each merge (a,b), find the token
|
||
* whose string == concat(a_str, b_str). This lets bpe.c skip O(n²) search. */
|
||
uint32_t *merge_result_ids = NULL;
|
||
if (n_merges > 0) {
|
||
merge_result_ids = malloc(n_merges * sizeof(uint32_t));
|
||
int resolved = 0;
|
||
for (int i = 0; i < n_merges; i++) {
|
||
/* Concatenate a + b merge strings */
|
||
const char *as = merge_a_str[i] ? merge_a_str[i] : "";
|
||
const char *bs = merge_b_str[i] ? merge_b_str[i] : "";
|
||
size_t alen = strlen(as), blen = strlen(bs);
|
||
char *merged = malloc(alen + blen + 1);
|
||
memcpy(merged, as, alen);
|
||
memcpy(merged + alen, bs, blen);
|
||
merged[alen + blen] = '\0';
|
||
int result_id = vocab_map_get(vmap, merged);
|
||
merge_result_ids[i] = (result_id >= 0) ? (uint32_t)result_id : 0;
|
||
if (result_id >= 0) resolved++;
|
||
free(merged);
|
||
}
|
||
fprintf(stderr, " Merge results resolved: %d/%d\n", resolved, n_merges);
|
||
}
|
||
|
||
/* Write merge IDs as two flat uint32 arrays + result IDs */
|
||
if (n_merges > 0) {
|
||
fwrite(merge_a_ids, sizeof(uint32_t), n_merges, out);
|
||
fwrite(merge_b_ids, sizeof(uint32_t), n_merges, out);
|
||
fwrite(merge_result_ids, sizeof(uint32_t), n_merges, out);
|
||
}
|
||
|
||
for (int i = 0; i < n_special; i++) {
|
||
uint32_t id = specials[i].id;
|
||
const char *name = specials[i].name;
|
||
uint16_t name_len = (uint16_t)strlen(name);
|
||
fwrite(&id, sizeof(uint32_t), 1, out);
|
||
fwrite(&name_len, sizeof(uint16_t), 1, out);
|
||
fwrite(name, 1, name_len, out);
|
||
}
|
||
|
||
/* Write byte_to_id[256]: maps raw byte → token ID using GPT-2 mapping.
|
||
* For each byte, find its GPT-2 unicode codepoint, look up the single-char
|
||
* token in vocab. bpe.c uses this for the initial byte→token step. */
|
||
uint32_t byte_to_id[256];
|
||
int byte_mapped = 0;
|
||
for (int b = 0; b < 256; b++) {
|
||
uint32_t cp = byte_to_unicode[b];
|
||
/* Encode codepoint to UTF-8 string for vocab lookup */
|
||
char utf8[5] = {0};
|
||
if (cp < 0x80) {
|
||
utf8[0] = (char)cp;
|
||
} else if (cp < 0x800) {
|
||
utf8[0] = (char)(0xC0 | (cp >> 6));
|
||
utf8[1] = (char)(0x80 | (cp & 0x3F));
|
||
} else {
|
||
utf8[0] = (char)(0xE0 | (cp >> 12));
|
||
utf8[1] = (char)(0x80 | ((cp >> 6) & 0x3F));
|
||
utf8[2] = (char)(0x80 | (cp & 0x3F));
|
||
}
|
||
int id = vocab_map_get(vmap, utf8);
|
||
byte_to_id[b] = (id >= 0) ? (uint32_t)id : 0;
|
||
if (id >= 0) byte_mapped++;
|
||
}
|
||
fprintf(stderr, " Byte-to-ID mapped: %d/256\n", byte_mapped);
|
||
fwrite(byte_to_id, sizeof(uint32_t), 256, out);
|
||
|
||
fclose(out);
|
||
|
||
long file_size = 0;
|
||
FILE *check = fopen(output_path, "rb");
|
||
if (check) { fseek(check, 0, SEEK_END); file_size = ftell(check); fclose(check); }
|
||
|
||
fprintf(stderr, " Vocab: %.1f KB -> %s\n",
|
||
(double)file_size / 1024.0, output_path);
|
||
|
||
cleanup_direct:
|
||
vocab_map_free(vmap);
|
||
for (int i = 0; i < n_tokens; i++) free(tokens[i]);
|
||
free(tokens);
|
||
for (int i = 0; i < n_tokens; i++) free(raw_tokens[i]);
|
||
free(raw_tokens); free(raw_lengths);
|
||
free(unicode_to_byte);
|
||
free(offsets); free(lengths); free(blob);
|
||
free(merge_a_ids); free(merge_b_ids); free(merge_result_ids);
|
||
if (merge_a_str) {
|
||
for (int i = 0; i < n_merges; i++) { free(merge_a_str[i]); free(merge_b_str[i]); }
|
||
free(merge_a_str); free(merge_b_str);
|
||
}
|
||
return true;
|
||
}
|
||
|
||
/* Fallback: use tokenizer.json via ft_load() */
|
||
if (tokenizer_path) {
|
||
fprintf(stderr, "Loading tokenizer: %s\n", tokenizer_path);
|
||
ft_tokenizer *tok = ft_load(tokenizer_path);
|
||
if (!tok) {
|
||
fprintf(stderr, "ERROR: failed to load tokenizer\n");
|
||
return false;
|
||
}
|
||
|
||
int n_tokens = ft_vocab_size(tok);
|
||
int n_merges = ft_num_merges(tok);
|
||
fprintf(stderr, " Vocab: %d tokens, %d merges\n", n_tokens, n_merges);
|
||
|
||
/* Build string blob */
|
||
size_t blob_size = 0;
|
||
for (int i = 0; i < n_tokens; i++) {
|
||
const char *s = ft_id_to_token(tok, i);
|
||
blob_size += s ? strlen(s) : 0;
|
||
}
|
||
|
||
uint32_t *offsets = malloc(n_tokens * sizeof(uint32_t));
|
||
uint16_t *lengths = malloc(n_tokens * sizeof(uint16_t));
|
||
char *blob = malloc(blob_size + 1);
|
||
size_t blob_pos = 0;
|
||
|
||
for (int i = 0; i < n_tokens; i++) {
|
||
const char *s = ft_id_to_token(tok, i);
|
||
size_t len = s ? strlen(s) : 0;
|
||
offsets[i] = (uint32_t)blob_pos;
|
||
lengths[i] = (uint16_t)len;
|
||
if (s && len > 0) {
|
||
memcpy(blob + blob_pos, s, len);
|
||
blob_pos += len;
|
||
}
|
||
}
|
||
|
||
/* Collect special tokens */
|
||
typedef struct { uint32_t id; const char *name; } special_entry_ft;
|
||
special_entry_ft specials[8];
|
||
int n_special = 0;
|
||
|
||
const char *special_names[] = {
|
||
"<|im_start|>", "<|im_end|>", "<|endoftext|>",
|
||
"<s>", "</s>", "<|begin_of_text|>", "<|end_of_text|>",
|
||
NULL
|
||
};
|
||
for (int i = 0; special_names[i] && n_special < 8; i++) {
|
||
int id = ft_token_to_id(tok, special_names[i]);
|
||
if (id >= 0) {
|
||
specials[n_special].id = (uint32_t)id;
|
||
specials[n_special].name = special_names[i];
|
||
n_special++;
|
||
}
|
||
}
|
||
|
||
fprintf(stderr, " Special tokens: %d\n", n_special);
|
||
|
||
FILE *out = fopen(output_path, "wb");
|
||
if (!out) {
|
||
fprintf(stderr, "ERROR: cannot create %s\n", output_path);
|
||
ft_free(tok);
|
||
free(offsets); free(lengths); free(blob);
|
||
return false;
|
||
}
|
||
|
||
btv_header hdr;
|
||
memset(&hdr, 0, sizeof(hdr));
|
||
memcpy(hdr.magic, "BTAV", 4);
|
||
hdr.n_tokens = n_tokens;
|
||
hdr.n_merges = n_merges;
|
||
hdr.n_special = n_special;
|
||
|
||
fwrite(&hdr, sizeof(hdr), 1, out);
|
||
fwrite(offsets, sizeof(uint32_t), n_tokens, out);
|
||
fwrite(lengths, sizeof(uint16_t), n_tokens, out);
|
||
fwrite(blob, 1, blob_pos, out);
|
||
|
||
/* Write merges as string pairs */
|
||
for (int i = 0; i < n_merges; i++) {
|
||
const char *a_str = NULL, *b_str = NULL;
|
||
ft_get_merge(tok, i, &a_str, &b_str);
|
||
uint16_t a_len = a_str ? (uint16_t)strlen(a_str) : 0;
|
||
uint16_t b_len = b_str ? (uint16_t)strlen(b_str) : 0;
|
||
fwrite(&a_len, sizeof(uint16_t), 1, out);
|
||
if (a_len) fwrite(a_str, 1, a_len, out);
|
||
fwrite(&b_len, sizeof(uint16_t), 1, out);
|
||
if (b_len) fwrite(b_str, 1, b_len, out);
|
||
}
|
||
|
||
for (int i = 0; i < n_special; i++) {
|
||
uint32_t id = specials[i].id;
|
||
const char *name = specials[i].name;
|
||
uint16_t name_len = (uint16_t)strlen(name);
|
||
fwrite(&id, sizeof(uint32_t), 1, out);
|
||
fwrite(&name_len, sizeof(uint16_t), 1, out);
|
||
fwrite(name, 1, name_len, out);
|
||
}
|
||
|
||
fclose(out);
|
||
|
||
long file_size = 0;
|
||
FILE *check = fopen(output_path, "rb");
|
||
if (check) { fseek(check, 0, SEEK_END); file_size = ftell(check); fclose(check); }
|
||
|
||
fprintf(stderr, " Vocab: %.1f KB -> %s\n",
|
||
(double)file_size / 1024.0, output_path);
|
||
|
||
ft_free(tok);
|
||
free(offsets); free(lengths); free(blob);
|
||
return true;
|
||
}
|
||
|
||
fprintf(stderr, "ERROR: no tokenizer files provided\n");
|
||
return false;
|
||
}
|