mini_embed 0.1.0 → 0.2.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: b44c5e93e9fc010a7c97e41f4fe4eaef3a5ee1033d1f8348c1d7a3f8d01d7d39
-  data.tar.gz: 61ca897bcf84b44822a15bc2ba9e37299567a797b26fd61ba987a4b7645c45b4
+  metadata.gz: 2df9f5c081f8a7fa2447261817ebba58e5c062921a1dc6ee3ec8048fdc300022
+  data.tar.gz: 4ee4f87161506c59e6deda7dd12b819c33e86ae5f5843aa89a9754b57b27f968
 SHA512:
-  metadata.gz: ab91e23951d0ef745d37a41467f778de5dcc7bedf21445292d31a1b6286add8cd851ebed6520756689d56226395bac3ce20a4f4bac0066cff8fc8488f26dcff6
-  data.tar.gz: 75c2885a6a63dfbc12db9bd80848f4d9526b7e37a23697ae227d1099a29ff5f6a3ea264378e5b66ce8a0b56935ddaa27de513594dfb530d5702652a09c391c27
+  metadata.gz: 4646b9f96a6ef525751d3046f0524d308940c1deb3a85623f775666ab2e1bbcddbec62c16f110f2cc1eae620f0b52c900a64f03cfa01f8e856d3546eb404ee98
+  data.tar.gz: f03b8103bddc296f1601d62bda851655529dc47cea1eee9f6bfcfc03fb5c005ad33f2b53a81e84d882b2813c670254bdbdb465b2b64f76c4ffc66248f5b27f73

data/README.md

@@ -1,74 +1,137 @@
 # mini_embed
 
-Fast, minimal GGUF embedding extractor for Ruby.
+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.
+
+---
+
+## Why MiniEmbed?
+
+- **Zero external dependencies** – no TensorFlow, PyTorch, or ONNX runtime.
+- **Single‑file C extension** – fast loading and mean‑pooled embeddings.
+- **Supports all common GGUF quantizations** – from `F32` to `Q2_K`.
+- **Works entirely offline** – your data never leaves your machine.
+- Perfect for **weekend projects**, **proof‑of‑concepts**, or **learning** about embeddings.
+
+---
 
 ## Installation
 
-Add to your Gemfile:
+Add this line to your application's `Gemfile`:
 
 ```ruby
 gem 'mini_embed'
 ```
-Or install globally:
 
-```sh
+Then execute:
+
+```bash
+bundle install
+```
+Or install it globally:
+
+```bash
 gem install mini_embed
 ```
 
-Usage
+
+## Requirements:
+
+A POSIX system (Linux, macOS, BSD) – Windows via WSL2 works.
+
+A C compiler and make (for compiling the native extension).
+
+A GGUF embedding model file (see Where to get models).
+
+## Usage
+
 ```ruby
 require 'mini_embed'
 
-model = MiniEmbed.new(model: 'path/to/model.gguf')
-embeddings_bin = model.embeddings(text: "hello world")  # => binary ouput
-embeddings_array = embeddings_bin.unpack('f*') # => array of float
-puts embeddings_array.size                     # => model dimension
-```
+# Load a GGUF model (F32, F16, Q8_0, Q4_K, etc. are all supported)
+model = MiniEmbed.new(model: '/path/to/gte-small.Q8_0.gguf')
 
-Supported Quantizations
+# Get embedding as an array of floats (default)
+embedding = model.embeddings(text: 'hello world')
+puts embedding.size   # e.g. 384
+puts embedding[0..4]  # e.g. [0.0123, -0.0456, ...]
 
+# Or get the raw binary string (little‑endian 32‑bit floats)
+binary = model.embeddings(text: 'hello world', type: :binary)
+embedding_from_binary = binary.unpack('e*')
 ```
-F32, F16
 
-Q4_0, Q4_1
+Note: The type parameter is optional – it defaults to :vector which returns a Ruby `Array<Float>`. Use `type: :binary` to get the raw binary string (compatible with the original C extension).
 
-Q5_0, Q5_1
 
-Q8_0, Q8_1
+## Simple tokenization note
+MiniEmbed uses a naive space‑based tokenizer. This means it splits input on spaces and looks up each token exactly in the model's vocabulary. For models trained with subword tokenization (like BERT), this will not work for out‑of‑vocabulary words.
+If you need proper subword tokenization, you can:
 
-Q2_K, Q3_K, Q4_K, Q5_K, Q6_K, Q8_K
-```
+- Pre‑tokenize in Ruby using the tokenizers gem and pass token IDs (not yet exposed in the C API, but easy to add).
+- Stick to simple vocabulary words that exist in the model (e.g., "text", "hello", "dog").
 
-## Building the Gem
+## Supported Quantization Types
 
-From the `mini_embed/` directory:
+| Type | Description   |
+|------|---------------|
+| 0    | F32 (float32) |
+| 1    | F16 (float16) |
+| 2    | Q4_0          |
+| 3    | Q4_1          |
+| 6    | Q5_0          |
+| 7    | Q5_1          |
+| 8    | Q8_0          |
+| 9    | Q8_1          |
+| 10   | Q2_K          |
+| 11   | Q3_K          |
+| 12   | Q4_K          |
+| 13   | Q5_K          |
+| 14   | Q6_K          |
+| 15   | Q8_K          |
 
-```bash
-bundle install
-bundle exec rake compile
-```
+The extension automatically dequantizes the embedding matrix on load, so inference speed is always that of a plain float32 lookup.
 
+Where to get models
+Hugging Face offers many GGUF models, e.g.:
 
-To build the gem file:
+- `gte-small`
+- `all‑MiniLM‑L6‑v2`
 
-```bash
-gem build mini_embed.gemspec
-```
+You can convert any safetensors or PyTorch model using the `convert‑hf‑to‑gguf.py` script from llama.cpp.
+
+For testing, we recommend the `gte-small` model (384 dimensions, ~30k vocabulary).
+
+## Limitations (Why this is not production‑ready)
+
+- Single‑threaded, blocking C code – embedding computation runs on the Ruby thread, freezing the interpreter.
+- No batching – only one text at a time.
+- Space‑based tokenization only – works only for words present exactly in the vocabulary.
+- Loads the entire embedding matrix into RAM – for large vocabularies this may consume significant memory.
+- No GPU support – CPU only.
+- Error handling is minimal – invalid models may crash the Ruby process.
 
-To install locally:
+If you need a robust, scalable solution, consider:
+
+- Running llama.cpp as a server (./server -m model.gguf --embeddings) and calling its HTTP endpoint.
+- Using a cloud embeddings API (OpenAI, Cohere, VoyageAI, etc.).
+- Deploying a dedicated inference service with BentoML or Ray Serve.
+
+
+## Development & Contributing
+Bug reports and pull requests are welcome on GitHub.
+To run the tests:
 
 ```bash
-gem install ./mini_embed-0.1.0.gem
+bundle exec rspec
 ```
-Using in a Rails project
-Add to Gemfile:
 
-```ruby
-gem 'mini_embed', path: '/path/to/mini_embed'
-```
+The gem uses rake-compiler to build the extension. After making changes to the C source, run:
 
-Then `bundle install` and use as above.
+```bash
+bundle exec rake compile
+```
 
 ## License
-
-MIT License. See [LICENSE](LICENSE).
+MIT License. See [LICENSE](LICENSE).

data/ext/mini_embed/mini_embed.c

@@ -13,6 +13,8 @@
 #define HASH_SIZE  131071
 #define MAX_DIMS   4
 #define GGUF_ALIGN 32
+#define MAX_MERGES 10000
+#define MAX_REGEX 256
 
 enum ggml_type {
     GGML_TYPE_F32  = 0,
@@ -31,7 +33,249 @@ enum ggml_type {
     GGML_TYPE_Q8_K = 15,
 };
 
+enum llama_vocab_type {
+    LLAMA_VOCAB_TYPE_NONE = 0,
+    LLAMA_VOCAB_TYPE_SPM  = 1,
+    LLAMA_VOCAB_TYPE_BPE  = 2,
+    LLAMA_VOCAB_TYPE_WPM  = 3,
+};
+
+/* ------------------------------------------------------------------------- */
+// Unicode helper functions
+static int unicode_len_utf8(char c) {
+    if ((c & 0x80) == 0) return 1;
+    if ((c & 0xE0) == 0xC0) return 2;
+    if ((c & 0xF0) == 0xE0) return 3;
+    if ((c & 0xF8) == 0xF0) return 4;
+    return 1;
+}
+
+static int unicode_is_letter(uint32_t cp) {
+    return (cp >= 0x41 && cp <= 0x5A) || (cp >= 0x61 && cp <= 0x7A) ||
+           (cp >= 0xC0 && cp <= 0xD6) || (cp >= 0xD8 && cp <= 0xF6) ||
+           (cp >= 0xF8 && cp <= 0x2FF) || (cp >= 0x370 && cp <= 0x37D) ||
+           (cp >= 0x37F && cp <= 0x1FFF) || (cp >= 0x200C && cp <= 0x200D) ||
+           (cp >= 0x2070 && cp <= 0x218F) || (cp >= 0x2C00 && cp <= 0x2FEF) ||
+           (cp >= 0x3001 && cp <= 0xD7FF) || (cp >= 0xF900 && cp <= 0xFDCF) ||
+           (cp >= 0xFDF0 && cp <= 0xFFFD);
+}
+
+static int unicode_is_number(uint32_t cp) {
+    return (cp >= 0x30 && cp <= 0x39) || (cp >= 0x660 && cp <= 0x669) ||
+           (cp >= 0x6F0 && cp <= 0x6F9) || (cp >= 0x7C0 && cp <= 0x7C9) ||
+           (cp >= 0x966 && cp <= 0x96F);
+}
+
+static uint32_t unicode_cpt_from_utf8(const char *s, size_t *len) {
+    unsigned char c = (unsigned char)s[0];
+    if (c < 0x80) { *len = 1; return c; }
+    if ((c & 0xE0) == 0xC0) { *len = 2; return ((c & 0x1F) << 6) | (s[1] & 0x3F); }
+    if ((c & 0xF0) == 0xE0) { *len = 3; return ((c & 0x0F) << 12) | ((s[1] & 0x3F) << 6) | (s[2] & 0x3F); }
+    if ((c & 0xF8) == 0xF0) { *len = 4; return ((c & 0x07) << 18) | ((s[1] & 0x3F) << 12) | ((s[2] & 0x3F) << 6) | (s[3] & 0x3F); }
+    *len = 1;
+    return c;
+}
+
+/* ------------------------------------------------------------------------- */
+// Simple regex pattern matcher (simplified)
+typedef struct {
+    char *pattern;
+    int pattern_len;
+} RegexPattern;
+
+static int match_regex(const char *text, const RegexPattern *patterns, int num_patterns) {
+    for (int i = 0; i < num_patterns; i++) {
+        const char *p = patterns[i].pattern;
+        if (strstr(p, "\\p{L}")) {
+            size_t len;
+            uint32_t cp = unicode_cpt_from_utf8(text, &len);
+            if (unicode_is_letter(cp)) return 1;
+        } else if (strstr(p, "\\p{N}")) {
+            size_t len;
+            uint32_t cp = unicode_cpt_from_utf8(text, &len);
+            if (unicode_is_number(cp)) return 1;
+        } else if (p[0] == '\\' && p[1] == 's') {
+            if (isspace(text[0])) return 1;
+        } else if (p[0] == '\\' && p[1] == 'r') {
+            if (text[0] == '\r') return 1;
+        } else if (p[0] == '\\' && p[1] == 'n') {
+            if (text[0] == '\n') return 1;
+        } else if (p[0] == '.' && p[1] == '*') {
+            return 1;
+        } else if (isalnum(p[0]) || ispunct(p[0])) {
+            if (text[0] == p[0]) return 1;
+        }
+    }
+    return 0;
+}
+
+static char** unicode_regex_split(const char *text, const RegexPattern *patterns, int num_patterns, int *num_words) {
+    char **words = NULL;
+    int word_count = 0, word_capacity = 0;
+    size_t text_len = strlen(text), pos = 0;
+
+    while (pos < text_len) {
+        size_t start = pos;
+        while (start < text_len && !match_regex(text + start, patterns, num_patterns)) start++;
+        if (start >= text_len) break;
+        size_t end = start;
+        while (end < text_len && match_regex(text + end, patterns, num_patterns)) end++;
+        if (end > start) {
+            size_t word_len = end - start;
+            char *word = malloc(word_len + 1);
+            if (!word) { while (--word_count >= 0) free(words[word_count]); free(words); *num_words = 0; return NULL; }
+            memcpy(word, text + start, word_len);
+            word[word_len] = '\0';
+            if (word_count >= word_capacity) {
+                word_capacity = word_capacity ? word_capacity * 2 : 16;
+                words = realloc(words, word_capacity * sizeof(char*));
+                if (!words) { free(word); while (--word_count >= 0) free(words[word_count]); *num_words = 0; return NULL; }
+            }
+            words[word_count++] = word;
+        }
+        pos = end;
+    }
+    *num_words = word_count;
+    return words;
+}
+
+/* ------------------------------------------------------------------------- */
+// BPE merge structures
+typedef struct {
+    char *left;
+    char *right;
+    char *merged;
+    int rank;
+} BPEMerge;
+
+typedef struct {
+    BPEMerge *merges;
+    int num_merges;
+    int capacity;
+} BPEMergeTable;
+
+static void bpe_merge_table_init(BPEMergeTable *table) {
+    memset(table, 0, sizeof(*table));
+}
+
+static void bpe_merge_table_add(BPEMergeTable *table, const char *left, const char *right, const char *merged, int rank) {
+    if (table->num_merges >= table->capacity) {
+        table->capacity = table->capacity ? table->capacity * 2 : 100;
+        table->merges = realloc(table->merges, table->capacity * sizeof(BPEMerge));
+    }
+    BPEMerge *m = &table->merges[table->num_merges++];
+    m->left = strdup(left);
+    m->right = strdup(right);
+    m->merged = strdup(merged);
+    m->rank = rank;
+}
+
+static void bpe_merge_table_free(BPEMergeTable *table) {
+    for (int i = 0; i < table->num_merges; i++) {
+        free(table->merges[i].left);
+        free(table->merges[i].right);
+        free(table->merges[i].merged);
+    }
+    free(table->merges);
+    table->merges = NULL;
+    table->num_merges = 0;
+}
+
+static int bpe_merge_rank(const BPEMergeTable *table, const char *left, const char *right) {
+    for (int i = 0; i < table->num_merges; i++) {
+        if (strcmp(table->merges[i].left, left) == 0 && strcmp(table->merges[i].right, right) == 0)
+            return table->merges[i].rank;
+    }
+    return -1;
+}
+
+/* ------------------------------------------------------------------------- */
+// BPE tokenization
+typedef struct {
+    char *text;
+    int start, end;
+    int prev, next;
+    int used;
+} BPESymbol;
+
+static void bpe_tokenize_word(const BPEMergeTable *merges, const char *word, int (*text_to_id)(void*, const char*), void *vocab_data, int *token_ids, int *num_tokens) {
+    int word_len = strlen(word);
+    int num_symbols = 0;
+    BPESymbol *symbols = malloc(word_len * sizeof(BPESymbol));
+    int offset = 0;
+    while (offset < word_len) {
+        int char_len = unicode_len_utf8(word[offset]);
+        symbols[num_symbols].text = (char*)word + offset;
+        symbols[num_symbols].start = offset;
+        symbols[num_symbols].end = offset + char_len;
+        symbols[num_symbols].prev = num_symbols - 1;
+        symbols[num_symbols].next = num_symbols + 1;
+        symbols[num_symbols].used = 1;
+        offset += char_len;
+        num_symbols++;
+    }
+
+    if (num_symbols <= 1) {
+        int id = text_to_id(vocab_data, word);
+        if (id != -1) token_ids[(*num_tokens)++] = id;
+        free(symbols);
+        return;
+    }
+
+    typedef struct { int left, right, rank; } Bigram;
+    Bigram *bigrams = malloc(num_symbols * num_symbols * sizeof(Bigram));
+    int num_bigrams = 0;
+    for (int i = 0; i < num_symbols - 1; i++) {
+        if (symbols[i].used && symbols[i+1].used) {
+            char *left_str = malloc(symbols[i].end - symbols[i].start + 1);
+            char *right_str = malloc(symbols[i+1].end - symbols[i+1].start + 1);
+            memcpy(left_str, symbols[i].text, symbols[i].end - symbols[i].start);
+            memcpy(right_str, symbols[i+1].text, symbols[i+1].end - symbols[i+1].start);
+            left_str[symbols[i].end - symbols[i].start] = '\0';
+            right_str[symbols[i+1].end - symbols[i+1].start] = '\0';
+            int rank = bpe_merge_rank(merges, left_str, right_str);
+            if (rank != -1) {
+                bigrams[num_bigrams].left = i;
+                bigrams[num_bigrams].right = i+1;
+                bigrams[num_bigrams].rank = rank;
+                num_bigrams++;
+            }
+            free(left_str); free(right_str);
+        }
+    }
+    for (int i = 0; i < num_bigrams - 1; i++)
+        for (int j = i+1; j < num_bigrams; j++)
+            if (bigrams[i].rank > bigrams[j].rank) {
+                Bigram tmp = bigrams[i];
+                bigrams[i] = bigrams[j];
+                bigrams[j] = tmp;
+            }
+
+    int *merged = calloc(num_symbols, sizeof(int));
+    for (int i = 0; i < num_bigrams; i++) {
+        int left = bigrams[i].left, right = bigrams[i].right;
+        if (merged[left] || merged[right]) continue;
+        symbols[left].end = symbols[right].end;
+        symbols[left].next = symbols[right].next;
+        merged[right] = 1;
+        if (symbols[right].next < num_symbols) symbols[symbols[right].next].prev = left;
+    }
+
+    for (int i = 0; i < num_symbols; i++) {
+        if (!merged[i] && symbols[i].used) {
+            char *substr = malloc(symbols[i].end - symbols[i].start + 1);
+            memcpy(substr, word + symbols[i].start, symbols[i].end - symbols[i].start);
+            substr[symbols[i].end - symbols[i].start] = '\0';
+            int id = text_to_id(vocab_data, substr);
+            if (id != -1) token_ids[(*num_tokens)++] = id;
+            free(substr);
+        }
+    }
+    free(bigrams); free(merged); free(symbols);
+}
+
 /* ------------------------------------------------------------------------- */
+// GGUF parsing
 static int safe_advance(uint8_t **p, uint8_t *end, size_t sz) {
     if (*p + sz > end) return 0;
     *p += sz;
@@ -39,14 +283,14 @@ static int safe_advance(uint8_t **p, uint8_t *end, size_t sz) {
 }
 
 static uint32_t rd32(uint8_t **p, uint8_t *end) {
-    uint32_t v = 0;
+    uint32_t v;
     if (!safe_advance(p, end, 4)) return 0;
     memcpy(&v, *p - 4, 4);
     return v;
 }
 
 static uint64_t rd64(uint8_t **p, uint8_t *end) {
-    uint64_t v = 0;
+    uint64_t v;
     if (!safe_advance(p, end, 8)) return 0;
     memcpy(&v, *p - 8, 8);
     return v;
@@ -57,9 +301,9 @@ static char *rdstr(uint8_t **p, uint8_t *end) {
     uint64_t len;
     memcpy(&len, *p, 8);
     *p += 8;
-    if (len == 0 || len > (1 << 20)) return NULL;
+    if (len == 0 || len > (1<<20)) return NULL;
     if (*p + len > end) return NULL;
-    char *s = malloc(len + 1);
+    char *s = malloc(len+1);
     if (!s) return NULL;
     memcpy(s, *p, len);
     s[len] = '\0';
@@ -70,27 +314,35 @@ static char *rdstr(uint8_t **p, uint8_t *end) {
 static void align_to_32(uint8_t **p, uint8_t *end, uint8_t *base) {
     size_t off = *p - base;
     size_t aligned = (off + GGUF_ALIGN - 1) & ~(GGUF_ALIGN - 1);
-    if (base + aligned <= end)
-        *p = base + aligned;
+    if (base + aligned <= end) *p = base + aligned;
 }
 
 /* ------------------------------------------------------------------------- */
+// Hash table for vocabulary
 typedef struct HashNode {
     char *key;
-    int   id;
+    int id;
     struct HashNode *next;
 } HashNode;
 
 typedef struct {
-    int        vocab_size;
-    int        dim;
-    char     **tokens;
-    float     *float_data;
-    void      *tensor_data;
-    int        tensor_type;
-    void      *mapped;
-    size_t     mapped_size;
+    int vocab_size;
+    int dim;
+    char **tokens;
+    float *float_data;
+    void *tensor_data;
+    int tensor_type;
+    void *mapped;
+    size_t mapped_size;
     HashNode **table;
+    BPEMergeTable merges;
+    RegexPattern *pre_patterns;
+    int num_pre_patterns;
+    int unknown_token_id;
+    int bos_token_id;
+    int eos_token_id;
+    int vocab_type;
+    char space_marker[8];
 } EmbedModel;
 
 typedef struct {
@@ -122,7 +374,12 @@ static int hget(EmbedModel *m, const char *k) {
     return -1;
 }
 
+static int text_to_id(void *vocab_data, const char *text) {
+    return hget((EmbedModel*)vocab_data, text);
+}
+
 /* ------------------------------------------------------------------------- */
+// File mapping
 static void *map_file(const char *path, size_t *size) {
     int fd = open(path, O_RDONLY);
     if (fd < 0) return NULL;
@@ -131,29 +388,22 @@ static void *map_file(const char *path, size_t *size) {
     *size = st.st_size;
     void *data = mmap(NULL, *size, PROT_READ, MAP_PRIVATE, fd, 0);
     close(fd);
-    if (data == MAP_FAILED) return NULL;
-    return data;
+    return data == MAP_FAILED ? NULL : data;
 }
 
 /* ------------------------------------------------------------------------- */
+// FP16 conversion
 static float fp16_to_fp32(uint16_t h) {
-    const uint16_t sign = (h >> 15) & 1;
-    const uint16_t exp  = (h >> 10) & 0x1F;
-    const uint16_t mant = h & 0x3FF;
-    float val;
-    if (exp == 0) {
-        val = (mant / 1024.0f) * 6.103515625e-5f;
-    } else if (exp == 31) {
-        return 0.0f;
-    } else {
-        val = (1.0f + mant / 1024.0f) * (1 << (exp - 15));
-    }
-    return sign ? -val : val;
+    uint16_t sign = (h >> 15) & 1;
+    uint16_t exp  = (h >> 10) & 0x1F;
+    uint16_t mant = h & 0x3FF;
+    if (exp == 0) return (mant / 1024.0f) * 6.103515625e-5f * (sign ? -1.0f : 1.0f);
+    if (exp == 31) return 0.0f;
+    return (1.0f + mant / 1024.0f) * (1 << (exp - 15)) * (sign ? -1.0f : 1.0f);
 }
 
 /* ------------------------------------------------------------------------- */
-/* Block dequantization */
-
+// Block dequantization functions
 static void dequantize_row_q4_0(const void *vx, float *y, int k) {
     const int nb = k / 32;
     const uint8_t *x = vx;
@@ -241,7 +491,6 @@ static void dequantize_row_q8_1(const void *vx, float *y, int k) {
     }
 }
 
-/* K-quants */
 static void dequantize_row_q2_K(const void *vx, float *y, int k) {
     const int nb = k / 256;
     const uint8_t *x = vx;
@@ -368,7 +617,6 @@ static void dequantize_row_q8_K(const void *vx, float *y, int k) {
     }
 }
 
-/* ------------------------------------------------------------------------- */
 static float* dequantize_tensor(const void *data, int type, int n_rows, int n_cols) {
     if (type == GGML_TYPE_F32) {
         float *out = malloc(n_rows * n_cols * sizeof(float));
@@ -413,6 +661,14 @@ static float* dequantize_tensor(const void *data, int type, int n_rows, int n_co
     for (int r = 0; r < n_rows; r++) {
         dequant_func(in + r * row_bytes, out + r * n_cols, n_cols);
     }
+
+    // Sanitize the tensor: replace NaNs, Infs, and astronomically large values with zero
+    int total = n_rows * n_cols;
+    for (int i = 0; i < total; i++) {
+        if (isnan(out[i]) || isinf(out[i]) || fabs(out[i]) > 1e10f) {
+            out[i] = 0.0f;
+        }
+    }
     return out;
 }
 
@@ -457,44 +713,89 @@ static void free_model_contents(EmbedModel *m) {
     }
     if (m->float_data) free(m->float_data);
     if (m->mapped) munmap(m->mapped, m->mapped_size);
+    bpe_merge_table_free(&m->merges);
+    if (m->pre_patterns) {
+        for (int i = 0; i < m->num_pre_patterns; i++) free(m->pre_patterns[i].pattern);
+        free(m->pre_patterns);
+    }
     free(m);
 }
 
 /* ------------------------------------------------------------------------- */
 static int is_printable_string(const char *s, size_t len) {
-    for (size_t i = 0; i < len; i++)
-        if (!isprint((unsigned char)s[i])) return 0;
+    for (size_t i = 0; i < len; i++) if (!isprint((unsigned char)s[i])) return 0;
     return 1;
 }
 
-/* Fallback: find the start of tensor info by scanning for a valid string */
 static uint8_t *find_tensor_info_start(uint8_t *cur, uint8_t *end) {
     uint8_t *scan = cur;
     while (scan + 8 < end) {
         uint64_t len;
         memcpy(&len, scan, 8);
-        if (len > 0 && len < 256 && scan + 8 + len <= end) {
-            if (is_printable_string((char*)scan + 8, len)) {
-                return scan;
-            }
-        }
+        if (len > 0 && len < 256 && scan + 8 + len <= end && is_printable_string((char*)scan+8, len))
+            return scan;
         scan++;
     }
     return NULL;
 }
 
+/* ------------------------------------------------------------------------- */
+static void detect_space_marker(EmbedModel *m) {
+    const char *candidates[] = {"▁", "Ġ", " "};
+    for (int i = 0; i < 3; i++) {
+        const char *marker = candidates[i];
+        int marker_len = strlen(marker);
+        for (int j = 0; j < m->vocab_size; j++) {
+            if (strncmp(m->tokens[j], marker, marker_len) == 0) {
+                strcpy(m->space_marker, marker);
+                return;
+            }
+        }
+    }
+    m->space_marker[0] = '\0';
+}
+
+static void setup_default_pre_patterns(EmbedModel *m) {
+    const char *default_patterns[] = {
+        "(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])",
+        "[^\\r\\n\\p{L}\\p{N}]?\\p{L}+",
+        "\\p{N}{1,3}",
+        " ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*",
+        "\\s*[\\r\\n]+",
+        "\\s+(?!\\S)",
+        "\\s+"
+    };
+    m->num_pre_patterns = sizeof(default_patterns)/sizeof(default_patterns[0]);
+    m->pre_patterns = malloc(m->num_pre_patterns * sizeof(RegexPattern));
+    for (int i = 0; i < m->num_pre_patterns; i++) {
+        m->pre_patterns[i].pattern = strdup(default_patterns[i]);
+        m->pre_patterns[i].pattern_len = strlen(default_patterns[i]);
+    }
+}
+
+static void parse_merge(const char *merge_str, char **left, char **right) {
+    const char *space = strchr(merge_str, ' ');
+    if (space) {
+        int left_len = space - merge_str;
+        *left = malloc(left_len+1);
+        memcpy(*left, merge_str, left_len);
+        (*left)[left_len] = '\0';
+        *right = strdup(space+1);
+    } else {
+        *left = strdup(merge_str);
+        *right = strdup("");
+    }
+}
+
 /* ------------------------------------------------------------------------- */
 static EmbedModel *embed_load_gguf(const char *path) {
     size_t sz;
     uint8_t *base = map_file(path, &sz);
     if (!base) return NULL;
-    uint8_t *cur = base;
-    uint8_t *end = base + sz;
-
+    uint8_t *cur = base, *end = base + sz;
     if (memcmp(cur, "GGUF", 4) != 0) { munmap(base, sz); return NULL; }
     cur += 4;
     uint32_t version = rd32(&cur, end);
-    (void)version;
     uint64_t n_tensors = rd64(&cur, end);
     uint64_t n_kv = rd64(&cur, end);
 
@@ -504,16 +805,20 @@ static EmbedModel *embed_load_gguf(const char *path) {
     m->mapped_size = sz;
     m->table = calloc(HASH_SIZE, sizeof(HashNode*));
     if (!m->table) { free_model_contents(m); return NULL; }
+    bpe_merge_table_init(&m->merges);
+    setup_default_pre_patterns(m);
+    m->unknown_token_id = -1;
+    m->bos_token_id = -1;
+    m->eos_token_id = -1;
+    m->vocab_type = LLAMA_VOCAB_TYPE_NONE;
+    m->space_marker[0] = '\0';
 
-    /* ---------- Metadata ---------- */
     int vocab_found = 0;
     for (uint64_t i = 0; i < n_kv; i++) {
         char *key = rdstr(&cur, end);
         if (!key) { free_model_contents(m); return NULL; }
         uint32_t type = rd32(&cur, end);
-
-        if ((strcmp(key, "tokenizer.ggml.tokens") == 0 ||
-             strcmp(key, "tokenizer.ggml.token_list") == 0) && type == 9) {
+        if ((strcmp(key, "tokenizer.ggml.tokens") == 0 || strcmp(key, "tokenizer.ggml.token_list") == 0) && type == 9) {
             uint32_t subtype = rd32(&cur, end);
             uint64_t n = rd64(&cur, end);
             if (subtype != 8) { free(key); free_model_contents(m); return NULL; }
@@ -527,21 +832,50 @@ static EmbedModel *embed_load_gguf(const char *path) {
                 hset(m, tok, (int)j);
             }
             vocab_found = 1;
-        } else {
-            if (!skip_value(&cur, end, type)) {
-                free(key); free_model_contents(m); return NULL;
+        } else if (strcmp(key, "tokenizer.ggml.merges") == 0 && type == 9) {
+            uint32_t subtype = rd32(&cur, end);
+            uint64_t n = rd64(&cur, end);
+            if (subtype == 8) {
+                for (uint64_t j = 0; j < n && j < MAX_MERGES; j++) {
+                    char *merge_str = rdstr(&cur, end);
+                    if (merge_str) {
+                        char *left, *right;
+                        parse_merge(merge_str, &left, &right);
+                        bpe_merge_table_add(&m->merges, left, right, merge_str, (int)j);
+                        free(left); free(right);
+                        free(merge_str);
+                    }
+                }
+            } else {
+                if (!skip_value(&cur, end, type)) { free(key); free_model_contents(m); return NULL; }
             }
+        } else if (strcmp(key, "tokenizer.ggml.model") == 0 && type == 8) {
+            char *model_type = rdstr(&cur, end);
+            if (model_type) {
+                if (strcmp(model_type, "gpt2") == 0 || strcmp(model_type, "llama") == 0) m->vocab_type = LLAMA_VOCAB_TYPE_BPE;
+                else if (strcmp(model_type, "bert") == 0) m->vocab_type = LLAMA_VOCAB_TYPE_WPM;
+                free(model_type);
+            }
+        } else if (strcmp(key, "tokenizer.ggml.pre") == 0 && type == 8) {
+            char *pre = rdstr(&cur, end);
+            if (pre) free(pre);
+        } else if (strcmp(key, "tokenizer.ggml.unknown_token_id") == 0 && type == 6) {
+            m->unknown_token_id = rd32(&cur, end);
+        } else if (strcmp(key, "tokenizer.ggml.bos_token_id") == 0 && type == 6) {
+            m->bos_token_id = rd32(&cur, end);
+        } else if (strcmp(key, "tokenizer.ggml.eos_token_id") == 0 && type == 6) {
+            m->eos_token_id = rd32(&cur, end);
+        } else {
+            if (!skip_value(&cur, end, type)) { free(key); free_model_contents(m); return NULL; }
         }
         free(key);
     }
-
     if (!vocab_found) { free_model_contents(m); return NULL; }
+    detect_space_marker(m);
 
     uint8_t *after_kv = cur;
     align_to_32(&cur, end, base);
     uint8_t *tensor_start = cur;
-
-    /* ---------- Tensor info ---------- */
     int embd_found = 0;
     void *raw_tensor_data = NULL;
     int tensor_type = -1;
@@ -555,39 +889,27 @@ static EmbedModel *embed_load_gguf(const char *path) {
             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);
+            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") != NULL)
-                    is_token_embd = 1;
-                else if ((uint64_t)m->vocab_size == dims[1] && strstr(name, "embd") != NULL)
-                    is_token_embd = 1;
+                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;
             }
-
             if (!embd_found && is_token_embd) {
                 if (n_dims < 2 || dims[1] == 0) { free(name); free_model_contents(m); return NULL; }
-                dim0 = dims[0];
-                dim1 = dims[1];
-                if (dim0 == (uint64_t)m->vocab_size) {
-                    m->dim = (int)dim1;
-                    need_transpose = 0;
-                } else if (dim1 == (uint64_t)m->vocab_size) {
-                    m->dim = (int)dim0;
-                    need_transpose = 1;
-                } else {
-                    m->dim = (dim0 < dim1) ? (int)dim0 : (int)dim1;
-                    need_transpose = (dim0 > dim1) ? 1 : 0;
-                }
+                dim0 = dims[0]; dim1 = dims[1];
+                if (dim0 == (uint64_t)m->vocab_size) { m->dim = (int)dim1; need_transpose = 0; }
+                else if (dim1 == (uint64_t)m->vocab_size) { m->dim = (int)dim0; need_transpose = 1; }
+                else { m->dim = (dim0 < dim1) ? (int)dim0 : (int)dim1; need_transpose = (dim0 > dim1) ? 1 : 0; }
                 raw_tensor_data = base + offset;
                 tensor_type = type;
                 embd_found = 1;
+                free(name);
+                break;
             }
             free(name);
         }
@@ -597,13 +919,8 @@ static EmbedModel *embed_load_gguf(const char *path) {
             if (!tensor_start) break;
         }
     }
+    if (!embd_found || m->dim == 0) { free_model_contents(m); return NULL; }
 
-    if (!embd_found || m->dim == 0) {
-        free_model_contents(m);
-        return NULL;
-    }
-
-    /* Dequantize */
     if (tensor_type == GGML_TYPE_F32 && !need_transpose) {
         m->float_data = NULL;
         m->tensor_data = raw_tensor_data;
@@ -611,10 +928,7 @@ static EmbedModel *embed_load_gguf(const char *path) {
         int n_rows = need_transpose ? (int)dim1 : (int)dim0;
         int n_cols = need_transpose ? (int)dim0 : (int)dim1;
         m->float_data = dequantize_tensor(raw_tensor_data, tensor_type, n_rows, n_cols);
-        if (!m->float_data) {
-            free_model_contents(m);
-            return NULL;
-        }
+        if (!m->float_data) { free_model_contents(m); return NULL; }
         m->tensor_data = m->float_data;
     }
     m->tensor_type = tensor_type;
@@ -625,36 +939,84 @@ static EmbedModel *embed_load_gguf(const char *path) {
 /* ------------------------------------------------------------------------- */
 static void embed_text(EmbedModel *m, const char *txt, float *out) {
     memset(out, 0, sizeof(float) * m->dim);
-    char *copy = strdup(txt);
-    if (!copy) return;
+    int num_words = 0;
+    char **words = unicode_regex_split(txt, m->pre_patterns, m->num_pre_patterns, &num_words);
+    if (!words || num_words == 0) {
+        // Fallback to simple space split
+        char *copy = strdup(txt);
+        if (copy) {
+            char *tok = strtok(copy, " \t\n\r");
+            int used = 0;
+            const float *embd = (float*)m->tensor_data;
+            while (tok) {
+                int id = hget(m, tok);
+                if (id >= 0 && id < m->vocab_size) {
+                    const float *vec = embd + id * m->dim;
+                    for (int i = 0; i < m->dim; i++) out[i] += vec[i];
+                    used++;
+                }
+                tok = strtok(NULL, " \t\n\r");
+            }
+            if (used) { float inv = 1.0f / used; for (int i = 0; i < m->dim; i++) out[i] *= inv; }
+            free(copy);
+        }
+        if (words) free(words);
+        return;
+    }
 
-    char *tok = strtok(copy, " ");
+    int *token_ids = malloc(m->vocab_size * sizeof(int));
     int used = 0;
-    const float *embd_matrix = m->tensor_data;
-
-    while (tok) {
-        int id = hget(m, tok);
-        if (id >= 0 && id < m->vocab_size) {
-            const float *vec = embd_matrix + id * m->dim;
-            for (int i = 0; i < m->dim; i++) out[i] += vec[i];
+    const float *embd = (float*)m->tensor_data;
+    for (int i = 0; i < num_words; i++) {
+        char *word = words[i];
+        int id = hget(m, word);
+        if (id == -1 && m->space_marker[0]) {
+            char *with_marker = malloc(strlen(m->space_marker) + strlen(word) + 1);
+            strcpy(with_marker, m->space_marker);
+            strcat(with_marker, word);
+            id = hget(m, with_marker);
+            free(with_marker);
+        }
+        if (id != -1) {
+            const float *vec = embd + id * m->dim;
+            for (int j = 0; j < m->dim; j++) out[j] += vec[j];
             used++;
+        } else {
+            int num_tokens = 0;
+            bpe_tokenize_word(&m->merges, word, text_to_id, m, token_ids, &num_tokens);
+            for (int k = 0; k < num_tokens; k++) {
+                int tid = token_ids[k];
+                if (tid >= 0 && tid < m->vocab_size) {
+                    const float *vec = embd + tid * m->dim;
+                    for (int j = 0; j < m->dim; j++) out[j] += vec[j];
+                    used++;
+                } else if (m->unknown_token_id != -1 && m->unknown_token_id < m->vocab_size) {
+                    const float *vec = embd + m->unknown_token_id * m->dim;
+                    for (int j = 0; j < m->dim; j++) out[j] += vec[j];
+                    used++;
+                }
+            }
         }
-        tok = strtok(NULL, " ");
+        free(word);
     }
-
+    free(words);
+    free(token_ids);
     if (used > 0) {
         float inv = 1.0f / used;
         for (int i = 0; i < m->dim; i++) out[i] *= inv;
     }
-    free(copy);
+    for (int i = 0; i < m->dim; i++) {
+        if (isnan(out[i]) || isinf(out[i])) {
+            out[i] = 0.0f;
+        }
+    }
 }
 
 /* ------------------------------------------------------------------------- */
+// Ruby bindings
 static void rb_embedder_free(void *p) {
     ruby_embedder *e = p;
-    if (!e) return;
-    if (e->model) free_model_contents(e->model);
-    free(e);
+    if (e) { if (e->model) free_model_contents(e->model); free(e); }
 }
 
 static size_t rb_embedder_memsize(const void *p) {
@@ -675,22 +1037,18 @@ static VALUE rb_embedder_alloc(VALUE klass) {
 static VALUE rb_embedder_initialize(VALUE self, VALUE opts) {
     ruby_embedder *e;
     TypedData_Get_Struct(self, ruby_embedder, &ruby_embedder_type, e);
-
     VALUE path = rb_hash_aref(opts, ID2SYM(rb_intern("model")));
     const char *cpath = StringValueCStr(path);
     e->model = embed_load_gguf(cpath);
-    if (!e->model)
-        rb_raise(rb_eRuntimeError, "failed to load GGUF model");
+    if (!e->model) rb_raise(rb_eRuntimeError, "failed to load GGUF model");
     return self;
 }
 
 static VALUE rb_embed(VALUE self, VALUE opts) {
     ruby_embedder *e;
     TypedData_Get_Struct(self, ruby_embedder, &ruby_embedder_type, e);
-
     VALUE text = rb_hash_aref(opts, ID2SYM(rb_intern("text")));
     const char *ctext = StringValueCStr(text);
-
     VALUE out = rb_str_new(NULL, e->model->dim * sizeof(float));
     embed_text(e->model, ctext, (float*)RSTRING_PTR(out));
     return out;
@@ -700,5 +1058,5 @@ void Init_mini_embed(void) {
     VALUE c = rb_define_class("MiniEmbed", rb_cObject);
     rb_define_alloc_func(c, rb_embedder_alloc);
     rb_define_method(c, "initialize", rb_embedder_initialize, 1);
-    rb_define_method(c, "embeddings", rb_embed, 1);
+    rb_define_method(c, "embed", rb_embed, 1);
 }

data/lib/mini_embed.rb

@@ -1,3 +1,17 @@
 # frozen_string_literal: true
 
 require 'mini_embed/mini_embed'
+
+class MiniEmbed
+  # @param text [String] - text to extract embeddings from
+  # @param type [Symbol, nil] - :binary or :vector - type of data you want to receive
+  # @return [String, <Float>] - type == :binary - binary string, type == :vector - array of floats
+  def embeddings(text: text, type: :vector)
+    binary_data = embed(text: text) # call original C method
+
+    return binary_data if type == :binary
+    return binary_data.unpack('e*') if type == :vector
+
+    raise ArgumentError, "Unsupported data type: #{type}"
+  end
+end

metadata

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