mini_embed 0.3.0 → 0.4.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: d6230ebfba3a401a8d26543f106e46952198bc6e89c5cf7632da40346933cf64
-  data.tar.gz: d5d37dd58c4bb3671053acb280db02ebb2ef78722d9c115f57f2594ad3a9ab50
+  metadata.gz: 0f3a2f9365c3ba228faf709ec97d986dadcc22a78ab0b706d35ba3da5e1552ce
+  data.tar.gz: 778847fe77dc4cb8b8774b62fe6f212c435880e470868f15c2517f1adb37211a
 SHA512:
-  metadata.gz: a826aad05808580120035f689412afdf976d77637cd9b1cb57df02740a7c86efff0120bf7fba172498e0b5d7ed82617bac99777e731d33e84bcbb823db543e29
-  data.tar.gz: f5bb3db889b9c51348daed59c3fbab9496237c3e9a64cb908ef386a1093e5e678531a5ad10eb051d0614dbe1fb9217d93a32049e6a5b8392b053d2474d6e9606
+  metadata.gz: '086948ced123967c0aa5f7e0fcb6624dffb2f68f6c95e4abd2dbaa429fb8717d98be76f6b173925b21af1bc14a0fe8af9d6d68d891ba4ce90d5a0b2145df55ef'
+  data.tar.gz: 2ac2c25baf87dd7b21fc38ccce6c3be0a3c133008ef27fa1790a95d3bc6146d5cb522166c93d8961edb16c3128e13d8f0d9206869abee43af10586f0124e00f2

data/README.md

@@ -1,5 +1,7 @@
 # mini_embed
 
+[![CircleCI](https://dl.circleci.com/status-badge/img/gh/Makapoxa/mini_embed/tree/main.svg?style=svg)](https://dl.circleci.com/status-badge/redirect/gh/Makapoxa/mini_embed/tree/main) [![Gem Version](https://badge.fury.io/rb/mini_embed.svg)](https://badge.fury.io/rb/mini_embed)
+
 A minimal, dependency‑free C extension for Ruby that loads [GGUF](https://github.com/ggerganov/ggml/blob/master/docs/gguf.md) embedding models and computes text embeddings **locally**.
 
 **⚠️ Important:** This gem is intended for **small projects, prototypes, and hobbyist use**. It allows you to experiment with embeddings without relying on external APIs or cloud costs. **Do not use MiniEmbed in production** – it lacks the performance, scalability, and tokenization robustness of dedicated solutions. For real applications, use a proper inference server like [llama.cpp](https://github.com/ggerganov/llama.cpp) with its HTTP API, or managed services such as OpenAI, Cohere, or Hugging Face.

data/ext/mini_embed/mini_embed.c

@@ -440,9 +440,23 @@ typedef struct HashNode {
     struct HashNode *next;
 } HashNode;
 
+typedef struct {
+    char *name;
+    uint32_t n_dims;
+    uint64_t dims[MAX_DIMS];
+    int type;
+    const uint8_t *data;
+    size_t row_bytes;
+} Tensor;
+
 typedef struct {
     int vocab_size;
     int dim;
+    int n_layers;
+    int n_heads;
+    int n_ctx;
+    int n_ff;
+    float eps;
     char **tokens;
     void *mapped;
     size_t mapped_size;
@@ -460,6 +474,11 @@ typedef struct {
     int need_transpose;
     uint64_t raw_dim0, raw_dim1;
     int normalize;
+    Tensor *tensors;
+    int n_tensors;
+    int sep_token_id;
+    int pad_token_id;
+    int cls_token_id;
 } EmbedModel;
 
 typedef struct {
@@ -541,6 +560,36 @@ static float fp16_to_fp32(uint16_t h) {
     return result;
 }
 
+static uint16_t fp32_to_fp16(float f) {
+    uint32_t x;
+    memcpy(&x, &f, sizeof(x));
+
+    uint32_t sign = (x >> 16) & 0x8000;
+    int exp = ((x >> 23) & 0xFF) - 127 + 15;
+    uint32_t mant = x & 0x7FFFFF;
+
+    if (exp <= 0) {
+        if (exp < -10) return (uint16_t)sign;
+        mant |= 0x800000;
+        uint32_t t = mant >> (1 - exp);
+        if (t & 0x00001000) t += 0x00002000;
+        return (uint16_t)(sign | (t >> 13));
+    } else if (exp >= 31) {
+        if (mant == 0) return (uint16_t)(sign | 0x7C00);
+        return (uint16_t)(sign | 0x7C00 | (mant >> 13));
+    } else {
+        if (mant & 0x00001000) {
+            mant += 0x00002000;
+            if (mant & 0x00800000) {
+                mant = 0;
+                exp += 1;
+            }
+        }
+        if (exp >= 31) return (uint16_t)(sign | 0x7C00);
+        return (uint16_t)(sign | ((uint32_t)exp << 10) | (mant >> 13));
+    }
+}
+
 /* ------------------------------------------------------------------------- */
 // Block dequantization functions (correct sizes)
 static void dequantize_row_q4_0(const void *vx, float *y, int k) {
@@ -552,9 +601,9 @@ static void dequantize_row_q4_0(const void *vx, float *y, int k) {
         memcpy(&d16, block, 2);
         const float d = fp16_to_fp32(d16);
         const uint8_t *q = block + 2;
-        for (int j = 0; j < 32; j++) {
-            const int v = (q[j/2] >> (4*(j%2))) & 0x0F;
-            y[i*32 + j] = (v - 8.0f) * d;
+        for (int j = 0; j < 16; j++) {
+            y[i*32 + j]      = ((q[j] & 0x0F) - 8.0f) * d;
+            y[i*32 + j + 16] = ((q[j] >> 4)   - 8.0f) * d;
         }
     }
 }
@@ -570,9 +619,9 @@ static void dequantize_row_q4_1(const void *vx, float *y, int k) {
         const float d = fp16_to_fp32(d16);
         const float m = fp16_to_fp32(m16);
         const uint8_t *q = block + 4;
-        for (int j = 0; j < 32; j++) {
-            const int v = (q[j/2] >> (4*(j%2))) & 0x0F;
-            y[i*32 + j] = v * d + m;
+        for (int j = 0; j < 16; j++) {
+            y[i*32 + j]      = (q[j] & 0x0F) * d + m;
+            y[i*32 + j + 16] = (q[j] >> 4)   * d + m;
         }
     }
 }
@@ -622,9 +671,10 @@ static void dequantize_row_q8_0(const void *vx, float *y, int k) {
     const uint8_t *x = vx;
     for (int i = 0; i < nb; i++) {
         const uint8_t *block = x + i * 34;
-        float d;
-        memcpy(&d, block, 4);
-        const int8_t *q = (const int8_t*)(block + 4);
+        uint16_t d16;
+        memcpy(&d16, block, 2);
+        const float d = fp16_to_fp32(d16);
+        const int8_t *q = (const int8_t*)(block + 2);
         for (int j = 0; j < 32; j++) {
             y[i*32 + j] = (float)q[j] * d;
         }
@@ -635,11 +685,13 @@ static void dequantize_row_q8_1(const void *vx, float *y, int k) {
     const int nb = k / QK8_0;
     const uint8_t *x = vx;
     for (int i = 0; i < nb; i++) {
-        const uint8_t *block = x + i * 40;
-        float d, s;
-        memcpy(&d, block, 4);
-        memcpy(&s, block + 4, 4);
-        const int8_t *q = (const int8_t*)(block + 8);
+        const uint8_t *block = x + i * 36;
+        uint16_t d16, s16;
+        memcpy(&d16, block, 2);
+        memcpy(&s16, block + 2, 2);
+        const float d = fp16_to_fp32(d16);
+        const float s = fp16_to_fp32(s16);
+        const int8_t *q = (const int8_t*)(block + 4);
         for (int j = 0; j < 32; j++) {
             y[i*32 + j] = (float)q[j] * d + s;
         }
@@ -932,7 +984,7 @@ static size_t get_row_bytes(int type, int n_cols) {
         case GGML_TYPE_Q5_0: return (n_cols / 32) * 22;
         case GGML_TYPE_Q5_1: return (n_cols / 32) * 24;
         case GGML_TYPE_Q8_0: return (n_cols / 32) * 34;
-        case GGML_TYPE_Q8_1: return (n_cols / 32) * 40;
+        case GGML_TYPE_Q8_1: return (n_cols / 32) * 36;
         case GGML_TYPE_Q2_K: return (n_cols / 256) * 84;
         case GGML_TYPE_Q3_K: return (n_cols / 256) * 110;
         case GGML_TYPE_Q4_K: return (n_cols / 256) * 144;
@@ -949,6 +1001,7 @@ static int skip_value(uint8_t **p, uint8_t *end, uint32_t type) {
         case 0: case 1: case 7: return safe_advance(p, end, 1);
         case 2: case 3:         return safe_advance(p, end, 2);
         case 4: case 5: case 6: return safe_advance(p, end, 4);
+        case 10: case 11: case 12: return safe_advance(p, end, 8);
         case 8: {
             uint64_t len = rd64(p, end);
             return safe_advance(p, end, len);
@@ -982,6 +1035,10 @@ static void free_model_contents(EmbedModel *m) {
         }
         free(m->table);
     }
+    if (m->tensors) {
+        for (int i = 0; i < m->n_tensors; i++) free(m->tensors[i].name);
+        free(m->tensors);
+    }
     if (m->mapped) munmap(m->mapped, m->mapped_size);
     bpe_merge_table_free(&m->merges);
     free(m);
@@ -1047,6 +1104,21 @@ static void parse_merge(const char *merge_str, char **left, char **right) {
     }
 }
 
+static Tensor *find_tensor(EmbedModel *m, const char *name) {
+    if (!m || !m->tensors) return NULL;
+    for (int i = 0; i < m->n_tensors; i++) {
+        if (strcmp(m->tensors[i].name, name) == 0) return &m->tensors[i];
+    }
+    return NULL;
+}
+
+static float rd_float32(uint8_t **p, uint8_t *end) {
+    uint32_t bits = rd32(p, end);
+    float v;
+    memcpy(&v, &bits, sizeof(v));
+    return v;
+}
+
 /* ------------------------------------------------------------------------- */
 static EmbedModel *embed_load_gguf(const char *path) {
     size_t sz;
@@ -1072,8 +1144,12 @@ static EmbedModel *embed_load_gguf(const char *path) {
     m->unknown_token_id = -1;
     m->bos_token_id = -1;
     m->eos_token_id = -1;
+    m->sep_token_id = -1;
+    m->pad_token_id = 0;
+    m->cls_token_id = -1;
     m->vocab_type = LLAMA_VOCAB_TYPE_NONE;
     m->normalize = NORM_NONE;
+    m->eps = 1e-12f;
 
     int vocab_found = 0;
     for (uint64_t i = 0; i < n_kv; i++) {
@@ -1136,12 +1212,31 @@ static EmbedModel *embed_load_gguf(const char *path) {
         } else if (strcmp(key, "tokenizer.ggml.pre") == 0 && type == 8) {
             char *pre = rdstr(&cur, end);
             free(pre);
-        } else if (strcmp(key, "tokenizer.ggml.unknown_token_id") == 0 && type == 6) {
+        } else if (strcmp(key, "bert.block_count") == 0 && type == 4) {
+            m->n_layers = (int)rd32(&cur, end);
+        } else if (strcmp(key, "bert.context_length") == 0 && type == 4) {
+            m->n_ctx = (int)rd32(&cur, end);
+        } else if (strcmp(key, "bert.embedding_length") == 0 && type == 4) {
+            m->dim = (int)rd32(&cur, end);
+        } else if (strcmp(key, "bert.feed_forward_length") == 0 && type == 4) {
+            m->n_ff = (int)rd32(&cur, end);
+        } else if (strcmp(key, "bert.attention.head_count") == 0 && type == 4) {
+            m->n_heads = (int)rd32(&cur, end);
+        } else if (strcmp(key, "bert.attention.layer_norm_epsilon") == 0 && type == 6) {
+            m->eps = rd_float32(&cur, end);
+        } else if (strcmp(key, "tokenizer.ggml.unknown_token_id") == 0 && type == 4) {
             m->unknown_token_id = (int)rd32(&cur, end);
-        } else if (strcmp(key, "tokenizer.ggml.bos_token_id") == 0 && type == 6) {
+        } else if (strcmp(key, "tokenizer.ggml.bos_token_id") == 0 && type == 4) {
             m->bos_token_id = (int)rd32(&cur, end);
-        } else if (strcmp(key, "tokenizer.ggml.eos_token_id") == 0 && type == 6) {
+        } else if (strcmp(key, "tokenizer.ggml.eos_token_id") == 0 && type == 4) {
             m->eos_token_id = (int)rd32(&cur, end);
+            m->sep_token_id = m->eos_token_id;
+        } else if (strcmp(key, "tokenizer.ggml.seperator_token_id") == 0 && type == 4) {
+            m->sep_token_id = (int)rd32(&cur, end);
+        } else if (strcmp(key, "tokenizer.ggml.padding_token_id") == 0 && type == 4) {
+            m->pad_token_id = (int)rd32(&cur, end);
+        } else if (strcmp(key, "tokenizer.ggml.cls_token_id") == 0 && type == 4) {
+            m->cls_token_id = (int)rd32(&cur, end);
         } else if (strcmp(key, "general.alignment") == 0 && type == 6) {
             rd32(&cur, end);
         } else {
@@ -1153,107 +1248,404 @@ static EmbedModel *embed_load_gguf(const char *path) {
     if (!vocab_found) { free_model_contents(m); return NULL; }
     detect_space_marker(m);
 
-    uint8_t *after_kv = cur;
+    m->tensors = calloc((size_t)n_tensors, sizeof(Tensor));
+    if (!m->tensors) { free_model_contents(m); return NULL; }
+    m->n_tensors = (int)n_tensors;
+
+    for (uint64_t i = 0; i < n_tensors; i++) {
+        Tensor *t = &m->tensors[i];
+        t->name = rdstr(&cur, end);
+        if (!t->name) { free_model_contents(m); return NULL; }
+        t->n_dims = rd32(&cur, end);
+        if (t->n_dims == 0 || t->n_dims > MAX_DIMS) { free_model_contents(m); return NULL; }
+        for (uint32_t d = 0; d < t->n_dims; d++) t->dims[d] = rd64(&cur, end);
+        t->type = (int)rd32(&cur, end);
+        uint64_t offset = rd64(&cur, end);
+        t->row_bytes = get_row_bytes(t->type, (int)t->dims[0]);
+        if (t->row_bytes == 0) { free_model_contents(m); return NULL; }
+        t->data = (const uint8_t*)(uintptr_t)offset;
+    }
+
     align_to_32(&cur, end, base);
-    uint8_t *tensor_start = cur;
-    int embd_found = 0;
-
-    for (int attempt = 0; attempt < 2; attempt++) {
-        cur = tensor_start;
-        for (uint64_t i = 0; i < n_tensors; i++) {
-            char *name = rdstr(&cur, end);
-            if (!name) break;
-            uint32_t n_dims = rd32(&cur, end);
-            uint64_t dims[MAX_DIMS] = {0};
-            for (uint32_t d = 0; d < n_dims && d < MAX_DIMS; d++) dims[d] = rd64(&cur, end);
-            uint32_t type   = rd32(&cur, end);
-            uint64_t offset = rd64(&cur, end);
-
-            int is_token_embd = (strcmp(name, "token_embd.weight") == 0 ||
-                                 strcmp(name, "embeddings.word_embeddings.weight") == 0 ||
-                                 strcmp(name, "model.embed_tokens.weight") == 0);
-
-            if (!is_token_embd && n_dims == 2 && m->vocab_size > 0) {
-                if ((uint64_t)m->vocab_size == dims[0] && strstr(name, "embd")) is_token_embd = 1;
-                else if ((uint64_t)m->vocab_size == dims[1] && strstr(name, "embd")) is_token_embd = 1;
+    uint8_t *data_start = cur;
+    for (int i = 0; i < m->n_tensors; i++) {
+        Tensor *t = &m->tensors[i];
+        uint64_t offset = (uint64_t)(uintptr_t)t->data;
+        size_t rows = t->n_dims > 1 ? (size_t)t->dims[1] : 1;
+        size_t total_size = rows * t->row_bytes;
+        if (offset > (uint64_t)sz || data_start + offset < data_start ||
+            data_start + offset + total_size > end) {
+            free_model_contents(m);
+            return NULL;
+        }
+        t->data = data_start + offset;
+    }
+
+    Tensor *embd = find_tensor(m, "token_embd.weight");
+    if (!embd) embd = find_tensor(m, "embeddings.word_embeddings.weight");
+    if (!embd || embd->n_dims < 2 || embd->dims[1] != (uint64_t)m->vocab_size) {
+        free_model_contents(m);
+        return NULL;
+    }
+
+    if (m->dim == 0) m->dim = (int)embd->dims[0];
+    if (m->n_ctx == 0) m->n_ctx = 512;
+    if (m->n_ff == 0) m->n_ff = m->dim * 4;
+    if (m->n_heads == 0) m->n_heads = 12;
+    if (m->n_layers == 0) m->n_layers = 12;
+    if (m->cls_token_id < 0) m->cls_token_id = m->bos_token_id;
+    if (m->sep_token_id < 0) m->sep_token_id = m->eos_token_id;
+
+    m->raw_tensor_data = embd->data;
+    m->tensor_type = embd->type;
+    m->row_bytes = embd->row_bytes;
+    m->raw_dim0 = embd->dims[0];
+    m->raw_dim1 = embd->dims[1];
+    m->need_transpose = 0;
+
+    if (m->dim <= 0 || m->dim > MAX_DIM) {
+        free_model_contents(m); return NULL; 
+    }
+
+    return m;
+}
+
+/* ------------------------------------------------------------------------- */
+// L2 normalization
+static void normalize_l2(float *vec, int dim) {
+    double sum = 0.0;
+    for (int i = 0; i < dim; i++) sum += vec[i] * vec[i];
+    double norm = sqrt(sum);
+    if (norm > 0.0) {
+        float inv = (float)(1.0 / norm);
+        for (int i = 0; i < dim; i++) vec[i] *= inv;
+    }
+}
+
+static void tensor_get_row(const Tensor *t, int row, float *out) {
+    if (!t || row < 0 || (t->n_dims > 1 && row >= (int)t->dims[1])) {
+        return;
+    }
+
+    const uint8_t *raw = t->data + (size_t)row * t->row_bytes;
+    int cols = (int)t->dims[0];
+    switch (t->type) {
+        case GGML_TYPE_F32:
+            memcpy(out, raw, (size_t)cols * sizeof(float));
+            break;
+        case GGML_TYPE_F16:
+            for (int i = 0; i < cols; i++) {
+                uint16_t h;
+                memcpy(&h, raw + (size_t)i * sizeof(uint16_t), sizeof(uint16_t));
+                out[i] = fp16_to_fp32(h);
             }
+            break;
+        case GGML_TYPE_Q4_0:
+            dequantize_row_q4_0(raw, out, cols);
+            break;
+        case GGML_TYPE_Q8_0:
+            dequantize_row_q8_0(raw, out, cols);
+            break;
+        default:
+            memset(out, 0, (size_t)cols * sizeof(float));
+            break;
+    }
+}
 
-            if (!embd_found && is_token_embd) {
-                if (n_dims < 2 || dims[1] == 0) { 
-                    free(name); free_model_contents(m); return NULL; 
-                }
-                
-                uint64_t ne0 = dims[0];
-                uint64_t ne1 = dims[1];
-
-                int need_transpose = 0;
-                int dim;
-                
-                if (ne1 == (uint64_t)m->vocab_size) {
-                    dim = (int)ne0;
-                    need_transpose = 0;
-                } else if (ne0 == (uint64_t)m->vocab_size) {
-                    dim = (int)ne1;
-                    need_transpose = 1;
-                } else {
-                    dim = (ne0 < ne1) ? (int)ne0 : (int)ne1;
-                    need_transpose = (ne0 > ne1) ? 1 : 0;
-                }
+static const float *tensor_f32_data(const Tensor *t) {
+    if (!t || t->type != GGML_TYPE_F32) return NULL;
+    return (const float*)t->data;
+}
 
-                if (dim <= 0 || dim > MAX_DIM) { 
-                    free(name); free_model_contents(m); return NULL; 
-                }
+static float dot_q4_0_q8_0_like_ggml(const uint8_t *raw, const float *x, int n) {
+    int nb = n / QK8_0;
+    float sumf = 0.0f;
 
-                size_t row_bytes = get_row_bytes(type, (int)(need_transpose ? ne1 : ne0));
-                size_t total_size = (size_t)(need_transpose ? ne1 : ne0) * row_bytes;
-                
-                if (offset >= sz || offset + total_size > sz) {
-                    free(name);
-                    free_model_contents(m);
-                    return NULL;
-                }
+    for (int ib = 0; ib < nb; ib++) {
+        const uint8_t *block = raw + (size_t)ib * 18;
+        uint16_t d16;
+        memcpy(&d16, block, 2);
+        const float dx = fp16_to_fp32(d16);
+        const uint8_t *q = block + 2;
+
+        const float *xb = x + (size_t)ib * QK8_0;
+        float amax = 0.0f;
+        for (int j = 0; j < QK8_0; j++) {
+            float av = fabsf(xb[j]);
+            if (av > amax) amax = av;
+        }
+
+        const float d = amax / 127.0f;
+        const float id = d ? 1.0f / d : 0.0f;
+        const float dy = fp16_to_fp32(fp32_to_fp16(d));
+        int8_t qy[QK8_0];
+        for (int j = 0; j < QK8_0; j++) qy[j] = (int8_t)roundf(xb[j] * id);
+
+        int sumi0 = 0;
+        int sumi1 = 0;
+        for (int j = 0; j < QK8_0/2; j++) {
+            const int v0 = (q[j] & 0x0F) - 8;
+            const int v1 = (q[j] >> 4) - 8;
+            sumi0 += v0 * qy[j];
+            sumi1 += v1 * qy[j + QK8_0/2];
+        }
+        sumf += (float)(sumi0 + sumi1) * dx * dy;
+    }
 
-                m->dim = dim;
-                m->raw_dim0 = ne0;
-                m->raw_dim1 = ne1;
-                m->need_transpose = need_transpose;
-                m->raw_tensor_data = base + offset;
-                m->tensor_type = type;
-                m->row_bytes = row_bytes;
-                embd_found = 1;
-                free(name);
+    return sumf;
+}
+
+static int ascii_wordpiece_tokenize(EmbedModel *m, const char *txt, int *ids, int max_ids) {
+    int n = 0;
+    if (m->cls_token_id >= 0 && n < max_ids) ids[n++] = m->cls_token_id;
+
+    size_t len = strlen(txt);
+    size_t i = 0;
+    while (i < len && n < max_ids - 1) {
+        while (i < len && isspace((unsigned char)txt[i])) i++;
+        if (i >= len) break;
+
+        char word[256];
+        int wl = 0;
+        if (isalnum((unsigned char)txt[i])) {
+            while (i < len && (isalnum((unsigned char)txt[i]) || txt[i] == '_') && wl < (int)sizeof(word) - 1) {
+                word[wl++] = (char)tolower((unsigned char)txt[i++]);
+            }
+            while (i < len && (isalnum((unsigned char)txt[i]) || txt[i] == '_')) i++;
+        } else {
+            word[wl++] = txt[i++];
+        }
+        word[wl] = '\0';
+        if (wl == 0) continue;
+
+        char word1[260];
+        const char marker[] = "\xE2\x96\x81";
+        memcpy(word1, marker, 3);
+        memcpy(word1 + 3, word, (size_t)wl + 1);
+        int w1l = wl + 3;
+
+        int current_tokens = n;
+        for (int start = 0; start < w1l && n < max_ids - 1; start++) {
+            int matched = 0;
+            for (int end_pos = w1l; end_pos > start; end_pos--) {
+                char piece[260];
+                int plen = end_pos - start;
+                memcpy(piece, word1 + start, plen);
+                piece[plen] = '\0';
+                int piece_id = hget(m, piece);
+                if (piece_id >= 0) {
+                    ids[n++] = piece_id;
+                    start = end_pos - 1;
+                    matched = 1;
+                    break;
+                }
+            }
+            if (!matched) {
+                n = current_tokens;
                 break;
             }
-            free(name);
         }
-        if (embd_found) break;
-        if (attempt == 0) {
-            tensor_start = find_tensor_info_start(after_kv, end);
-            if (!tensor_start) break;
+
+        if (n == current_tokens && m->unknown_token_id >= 0 && n < max_ids - 1) ids[n++] = m->unknown_token_id;
+    }
+
+    if (m->sep_token_id >= 0 && n < max_ids) ids[n++] = m->sep_token_id;
+    return n;
+}
+
+static void linear_one(const Tensor *w, const Tensor *b, const float *x, float *out, float *row) {
+    int in = (int)w->dims[0];
+    int out_dim = (int)w->dims[1];
+    const float *bias = tensor_f32_data(b);
+    for (int o = 0; o < out_dim; o++) {
+        float sum = bias ? bias[o] : 0.0f;
+        if (w->type == GGML_TYPE_Q4_0) {
+            const uint8_t *raw = w->data + (size_t)o * w->row_bytes;
+            sum += dot_q4_0_q8_0_like_ggml(raw, x, in);
+        } else {
+            tensor_get_row(w, o, row);
+            for (int i = 0; i < in; i++) sum += row[i] * x[i];
         }
+        out[o] = sum;
     }
+}
 
-    if (!embd_found || m->dim == 0) { 
-        free_model_contents(m); return NULL; 
+static void linear_batch(const Tensor *w, const Tensor *b, const float *x, int seq, float *out, float *row) {
+    int in = (int)w->dims[0];
+    int out_dim = (int)w->dims[1];
+    for (int t = 0; t < seq; t++) {
+        linear_one(w, b, x + (size_t)t * in, out + (size_t)t * out_dim, row);
     }
+}
 
-    return m;
+static void layer_norm(const float *x, const Tensor *w, const Tensor *b, int seq, int dim, float eps, float *out) {
+    const float *weight = tensor_f32_data(w);
+    const float *bias = tensor_f32_data(b);
+    for (int t = 0; t < seq; t++) {
+        const float *src = x + (size_t)t * dim;
+        float *dst = out + (size_t)t * dim;
+        float mean = 0.0f;
+        for (int i = 0; i < dim; i++) mean += src[i];
+        mean /= (float)dim;
+        float var = 0.0f;
+        for (int i = 0; i < dim; i++) {
+            float d = src[i] - mean;
+            var += d * d;
+        }
+        var /= (float)dim;
+        float scale = 1.0f / sqrtf(var + eps);
+        for (int i = 0; i < dim; i++) {
+            dst[i] = (src[i] - mean) * scale * (weight ? weight[i] : 1.0f) + (bias ? bias[i] : 0.0f);
+        }
+    }
 }
 
-/* ------------------------------------------------------------------------- */
-// L2 normalization
-static void normalize_l2(float *vec, int dim) {
-    float sum = 0;
-    for (int i = 0; i < dim; i++) sum += vec[i] * vec[i];
-    float norm = sqrtf(sum);
-    if (norm > 1e-8f) {
-        float inv = 1.0f / norm;
-        for (int i = 0; i < dim; i++) vec[i] *= inv;
+static float gelu_approx(float x) {
+    if (x <= -10.0f) return 0.0f;
+    if (x >= 10.0f) return x;
+    const float c = 0.7978845608028654f;
+    float hx = fp16_to_fp32(fp32_to_fp16(x));
+    float y = 0.5f * hx * (1.0f + tanhf(c * hx * (1.0f + 0.044715f * hx * hx)));
+    return fp16_to_fp32(fp32_to_fp16(y));
+}
+
+static int bert_embed_text(EmbedModel *m, const char *txt, float *out) {
+    if (m->vocab_type != LLAMA_VOCAB_TYPE_WPM || !find_tensor(m, "blk.0.attn_q.weight")) return 0;
+
+    memset(out, 0, (size_t)m->dim * sizeof(float));
+    if (!txt || !*txt) return 1;
+
+    int max_seq = m->n_ctx > 0 ? m->n_ctx : 512;
+    int *ids = malloc((size_t)max_seq * sizeof(int));
+    if (!ids) return 1;
+    int seq = ascii_wordpiece_tokenize(m, txt, ids, max_seq);
+    if (seq <= 0) { free(ids); return 1; }
+
+    int dim = m->dim;
+    int ff = m->n_ff;
+    int heads = m->n_heads;
+    int head_dim = dim / heads;
+    float *hidden = calloc((size_t)seq * dim, sizeof(float));
+    float *tmp = calloc((size_t)seq * dim, sizeof(float));
+    float *q = calloc((size_t)seq * dim, sizeof(float));
+    float *k = calloc((size_t)seq * dim, sizeof(float));
+    float *v = calloc((size_t)seq * dim, sizeof(float));
+    float *ctx = calloc((size_t)seq * dim, sizeof(float));
+    float *proj = calloc((size_t)seq * dim, sizeof(float));
+    float *ffn = calloc((size_t)seq * ff, sizeof(float));
+    float *row = malloc((size_t)(ff > dim ? ff : dim) * sizeof(float));
+    float *scores = malloc((size_t)seq * sizeof(float));
+    if (!hidden || !tmp || !q || !k || !v || !ctx || !proj || !ffn || !row || !scores) {
+        free(ids); free(hidden); free(tmp); free(q); free(k); free(v); free(ctx); free(proj); free(ffn); free(row); free(scores);
+        return 1;
+    }
+
+    Tensor *tok_emb = find_tensor(m, "token_embd.weight");
+    Tensor *pos_emb = find_tensor(m, "position_embd.weight");
+    Tensor *typ_emb = find_tensor(m, "token_types.weight");
+    Tensor *emb_norm_w = find_tensor(m, "token_embd_norm.weight");
+    Tensor *emb_norm_b = find_tensor(m, "token_embd_norm.bias");
+
+    float *tok = row;
+    float *pos = malloc((size_t)dim * sizeof(float));
+    float *typ = malloc((size_t)dim * sizeof(float));
+    if (!tok_emb || !pos_emb || !typ_emb || !pos || !typ) {
+        free(ids); free(hidden); free(tmp); free(q); free(k); free(v); free(ctx); free(proj); free(ffn); free(row); free(scores); free(pos); free(typ);
+        return 1;
     }
+
+    for (int t = 0; t < seq; t++) {
+        tensor_get_row(tok_emb, ids[t], tok);
+        tensor_get_row(pos_emb, t, pos);
+        tensor_get_row(typ_emb, 0, typ);
+        for (int d = 0; d < dim; d++) hidden[(size_t)t * dim + d] = tok[d] + pos[d] + typ[d];
+    }
+    layer_norm(hidden, emb_norm_w, emb_norm_b, seq, dim, m->eps, tmp);
+    memcpy(hidden, tmp, (size_t)seq * dim * sizeof(float));
+
+    for (int layer = 0; layer < m->n_layers; layer++) {
+        char name[80];
+        #define TENSOR(suffix) (snprintf(name, sizeof(name), "blk.%d.%s", layer, suffix), find_tensor(m, name))
+        Tensor *qw = TENSOR("attn_q.weight");
+        Tensor *qb = TENSOR("attn_q.bias");
+        Tensor *kw = TENSOR("attn_k.weight");
+        Tensor *kb = TENSOR("attn_k.bias");
+        Tensor *vw = TENSOR("attn_v.weight");
+        Tensor *vb = TENSOR("attn_v.bias");
+        Tensor *ow = TENSOR("attn_output.weight");
+        Tensor *ob = TENSOR("attn_output.bias");
+        Tensor *an_w = TENSOR("attn_output_norm.weight");
+        Tensor *an_b = TENSOR("attn_output_norm.bias");
+        Tensor *fu_w = TENSOR("ffn_up.weight");
+        Tensor *fu_b = TENSOR("ffn_up.bias");
+        Tensor *fd_w = TENSOR("ffn_down.weight");
+        Tensor *fd_b = TENSOR("ffn_down.bias");
+        Tensor *ln_w = TENSOR("layer_output_norm.weight");
+        Tensor *ln_b = TENSOR("layer_output_norm.bias");
+        #undef TENSOR
+
+        if (!qw || !qb || !kw || !kb || !vw || !vb || !ow || !ob || !an_w || !an_b ||
+            !fu_w || !fu_b || !fd_w || !fd_b || !ln_w || !ln_b) break;
+
+        linear_batch(qw, qb, hidden, seq, q, row);
+        linear_batch(kw, kb, hidden, seq, k, row);
+        linear_batch(vw, vb, hidden, seq, v, row);
+        memset(ctx, 0, (size_t)seq * dim * sizeof(float));
+
+        float att_scale = 1.0f / sqrtf((float)head_dim);
+        for (int h = 0; h < heads; h++) {
+            int off = h * head_dim;
+            for (int ti = 0; ti < seq; ti++) {
+                float max_score = -INFINITY;
+                for (int tj = 0; tj < seq; tj++) {
+                    float dot = 0.0f;
+                    const float *qv0 = q + (size_t)ti * dim + off;
+                    const float *kv0 = k + (size_t)tj * dim + off;
+                    for (int d = 0; d < head_dim; d++) dot += qv0[d] * kv0[d];
+                    scores[tj] = dot * att_scale;
+                    if (scores[tj] > max_score) max_score = scores[tj];
+                }
+                double sum = 0.0;
+                for (int tj = 0; tj < seq; tj++) {
+                    scores[tj] = expf(scores[tj] - max_score);
+                    sum += scores[tj];
+                }
+                float inv_sum = (float)(1.0 / sum);
+                float *dst = ctx + (size_t)ti * dim + off;
+                for (int tj = 0; tj < seq; tj++) {
+                    float p = scores[tj] * inv_sum;
+                    const float *vv0 = v + (size_t)tj * dim + off;
+                    for (int d = 0; d < head_dim; d++) dst[d] += p * vv0[d];
+                }
+            }
+        }
+
+        linear_batch(ow, ob, ctx, seq, proj, row);
+        for (int i = 0; i < seq * dim; i++) tmp[i] = hidden[i] + proj[i];
+        layer_norm(tmp, an_w, an_b, seq, dim, m->eps, hidden);
+
+        linear_batch(fu_w, fu_b, hidden, seq, ffn, row);
+        for (int i = 0; i < seq * ff; i++) ffn[i] = gelu_approx(ffn[i]);
+        linear_batch(fd_w, fd_b, ffn, seq, proj, row);
+        for (int i = 0; i < seq * dim; i++) tmp[i] = hidden[i] + proj[i];
+        layer_norm(tmp, ln_w, ln_b, seq, dim, m->eps, hidden);
+    }
+
+    for (int t = 0; t < seq; t++) {
+        for (int d = 0; d < dim; d++) out[d] += hidden[(size_t)t * dim + d];
+    }
+    float inv = 1.0f / (float)seq;
+    for (int d = 0; d < dim; d++) out[d] *= inv;
+    normalize_l2(out, dim);
+
+    free(ids); free(hidden); free(tmp); free(q); free(k); free(v); free(ctx); free(proj); free(ffn); free(row); free(scores); free(pos); free(typ);
+    return 1;
 }
 
 /* ------------------------------------------------------------------------- */
 static void embed_text(EmbedModel *m, const char *txt, float *out) {
+    if (bert_embed_text(m, txt, out)) return;
+
     memset(out, 0, sizeof(float) * m->dim);
     if (!txt || !*txt) return;
 
@@ -1413,4 +1805,4 @@ void Init_mini_embed(void) {
     rb_define_alloc_func(c, rb_embedder_alloc);
     rb_define_method(c, "initialize", rb_embedder_initialize, 1);
     rb_define_method(c, "embed", rb_embed, 1);
-}
+}

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: mini_embed
 version: !ruby/object:Gem::Version
-  version: 0.3.0
+  version: 0.4.0
 platform: ruby
 authors:
 - Makapoxa