static_embeddings 0.1.2 → 0.1.4

data/ext/static_embeddings/static_embeddings.c

@@ -274,10 +274,14 @@ static void batch_worker_run(batch_job_t *job, se_scratch_t *scratch) {
 
 static void *batch_execute(void *arg) {
     batch_job_t *job = (batch_job_t *)arg;
-    se_scratch_t scratch;
-    se_scratch_init(&scratch);
-    batch_worker_run(job, &scratch);
-    se_scratch_free(&scratch);
+    se_scratch_t *scratch = se_scratch_acquire(job->model->meta.dim);
+    if (!scratch) {
+        se_error_set(&job->error, SE_ERR_OOM, "out of memory while sizing scratch buffers");
+        job->failed = 1;
+        return NULL;
+    }
+    batch_worker_run(job, scratch);
+    se_scratch_release(scratch);
     return NULL;
 }
 
@@ -503,6 +507,51 @@ static VALUE lookup_option(VALUE opts, ID id) {
     return rb_hash_lookup2(opts, ID2SYM(id), Qundef);
 }
 
+typedef struct {
+    const ID *allowed;
+    size_t count;
+} keyword_check_t;
+
+static int reject_unknown_keyword_i(VALUE key, VALUE value, VALUE arg) {
+    const keyword_check_t *check = (const keyword_check_t *)arg;
+    (void)value;
+
+    if (!SYMBOL_P(key))
+        rb_raise(rb_eArgError, "keyword must be a Symbol");
+
+    ID id = SYM2ID(key);
+    for (size_t i = 0; i < check->count; i++) {
+        if (id == check->allowed[i])
+            return ST_CONTINUE;
+    }
+
+    VALUE names = rb_ary_new_capa((long)check->count);
+    for (size_t i = 0; i < check->count; i++)
+        rb_ary_push(names, rb_sprintf(":%s", rb_id2name(check->allowed[i])));
+
+    const char *name = rb_id2name(id);
+    rb_raise(rb_eArgError, "unknown keyword: :%s (accepted: %" PRIsVALUE ")", name ? name : "?",
+             rb_ary_join(names, rb_str_new_cstr(", ")));
+}
+
+static void check_keywords(VALUE opts, const ID *allowed, size_t count) {
+    if (NIL_P(opts))
+        return;
+    if (!RB_TYPE_P(opts, T_HASH))
+        rb_raise(rb_eArgError, "keywords must be a Hash");
+
+    keyword_check_t check;
+    check.allowed = allowed;
+    check.count = count;
+    rb_hash_foreach(opts, reject_unknown_keyword_i, (VALUE)&check);
+}
+
+#define SE_CHECK_KEYWORDS(opts, ...)                                                       \
+    do {                                                                                   \
+        const ID se_allowed_[] = {__VA_ARGS__};                                            \
+        check_keywords((opts), se_allowed_, sizeof(se_allowed_) / sizeof(se_allowed_[0])); \
+    } while (0)
+
 static void reject_parallel_threads(VALUE opts) {
     VALUE v = lookup_option(opts, id_threads);
     if (v == Qundef || v == Qnil)
@@ -867,6 +916,7 @@ static VALUE embed_batch_internal(VALUE self, VALUE texts, VALUE max_tokens_opt,
 static VALUE model_embed_batch(int argc, VALUE *argv, VALUE self) {
     VALUE texts, opts;
     rb_scan_args(argc, argv, "1:", &texts, &opts);
+    SE_CHECK_KEYWORDS(opts, id_max_tokens, id_format, id_validate_encoding, id_threads);
     reject_parallel_threads(opts);
 
     VALUE max_tokens = lookup_option(opts, id_max_tokens);
@@ -882,6 +932,32 @@ static VALUE embed_one_via_batch(VALUE self, VALUE text, VALUE max_tokens_opt,
     return embed_texts_internal(self, text, 0, 1, max_tokens_opt, format, validation, stats);
 }
 
+typedef struct {
+    se_scratch_t *scratch;
+    const se_model_t *model;
+    const uint8_t *input;
+    size_t input_len;
+    uint32_t max_tokens;
+    float *out;
+    se_token_stats_t *stats;
+    se_error_t err;
+    se_status_t rc;
+} embed_one_scratch_job_t;
+
+static VALUE embed_one_scratch_body(VALUE arg) {
+    embed_one_scratch_job_t *job = (embed_one_scratch_job_t *)(uintptr_t)arg;
+    job->rc = se_embed_one(job->model, job->scratch, job->input, job->input_len, job->max_tokens,
+                           job->out, job->stats, &job->err, NULL);
+    return Qnil;
+}
+
+static VALUE embed_one_scratch_ensure(VALUE arg) {
+    embed_one_scratch_job_t *job = (embed_one_scratch_job_t *)(uintptr_t)arg;
+    se_scratch_release(job->scratch);
+    job->scratch = NULL;
+    return Qnil;
+}
+
 static VALUE embed_one_value(VALUE self, VALUE text, VALUE max_tokens_opt,
                              se_vector_format_t format, se_encoding_validation_t validation,
                              se_token_stats_t *stats) {
@@ -901,25 +977,28 @@ static VALUE embed_one_value(VALUE self, VALUE text, VALUE max_tokens_opt,
     rb_enc_associate(result, binary_encoding);
     float *out = (float *)RSTRING_PTR(result);
 
-    se_scratch_t scratch;
-    se_scratch_init(&scratch);
-    if (!se_scratch_reserve(&scratch, dim)) {
-        se_scratch_free(&scratch);
+    se_token_stats_t local_stats;
+    embed_one_scratch_job_t job;
+    memset(&job, 0, sizeof(job));
+    job.model = &w->model;
+    job.input = (const uint8_t *)RSTRING_PTR(text);
+    job.input_len = (size_t)RSTRING_LEN(text);
+    job.max_tokens = resolve_max_tokens(&w->model, max_tokens_opt);
+    job.out = out;
+    job.stats = stats ? stats : &local_stats;
+    se_error_clear(&job.err);
+
+    job.scratch = se_scratch_acquire(dim);
+    if (!job.scratch)
         rb_raise(rb_eNoMemError, "out of memory");
-    }
 
-    se_error_t err;
-    se_error_clear(&err);
-    se_token_stats_t local_stats;
-    se_status_t rc =
-        se_embed_one(&w->model, &scratch, (const uint8_t *)RSTRING_PTR(text),
-                     (size_t)RSTRING_LEN(text), resolve_max_tokens(&w->model, max_tokens_opt), out,
-                     stats ? stats : &local_stats, &err, NULL);
-    se_scratch_free(&scratch);
+    rb_ensure(embed_one_scratch_body, (VALUE)(uintptr_t)&job, embed_one_scratch_ensure,
+              (VALUE)(uintptr_t)&job);
     RB_GC_GUARD(text);
+    RB_GC_GUARD(result);
 
-    if (rc != SE_OK)
-        raise_se(&err);
+    if (job.rc != SE_OK)
+        raise_se(&job.err);
 
     return result;
 }
@@ -927,6 +1006,7 @@ static VALUE embed_one_value(VALUE self, VALUE text, VALUE max_tokens_opt,
 static VALUE model_embed(int argc, VALUE *argv, VALUE self) {
     VALUE text, opts;
     rb_scan_args(argc, argv, "1:", &text, &opts);
+    SE_CHECK_KEYWORDS(opts, id_max_tokens, id_format, id_validate_encoding, id_threads);
     reject_parallel_threads(opts);
     return embed_one_value(self, text, lookup_option(opts, id_max_tokens),
                            resolve_vector_format(lookup_option(opts, id_format)),
@@ -937,6 +1017,7 @@ static VALUE model_embed(int argc, VALUE *argv, VALUE self) {
 static VALUE model_embed_with_stats(int argc, VALUE *argv, VALUE self) {
     VALUE text, opts;
     rb_scan_args(argc, argv, "1:", &text, &opts);
+    SE_CHECK_KEYWORDS(opts, id_max_tokens, id_format, id_validate_encoding, id_threads);
     reject_parallel_threads(opts);
 
     se_token_stats_t stats;
@@ -953,41 +1034,69 @@ static VALUE model_embed_with_stats(int argc, VALUE *argv, VALUE self) {
     return hash;
 }
 
+typedef struct {
+    se_scratch_t *scratch;
+    const se_model_t *model;
+    const uint8_t *input;
+    size_t input_len;
+    uint32_t max_tokens;
+    se_token_stats_t stats;
+    se_error_t err;
+    se_status_t rc;
+    VALUE ids;
+} tokenize_scratch_job_t;
+
+static VALUE tokenize_scratch_body(VALUE arg) {
+    tokenize_scratch_job_t *job = (tokenize_scratch_job_t *)(uintptr_t)arg;
+    job->rc = se_tokenize(job->model, job->scratch, job->input, job->input_len, job->max_tokens,
+                          &job->stats, &job->err, NULL);
+    if (job->rc != SE_OK)
+        return Qnil;
+
+    job->ids = rb_ary_new_capa((long)job->stats.token_count);
+    for (uint32_t i = 0; i < job->stats.token_count; i++)
+        rb_ary_push(job->ids, UINT2NUM(job->scratch->ids[i]));
+    return Qnil;
+}
+
+static VALUE tokenize_scratch_ensure(VALUE arg) {
+    tokenize_scratch_job_t *job = (tokenize_scratch_job_t *)(uintptr_t)arg;
+    se_scratch_release(job->scratch);
+    job->scratch = NULL;
+    return Qnil;
+}
+
 static VALUE model_tokenize(int argc, VALUE *argv, VALUE self) {
     VALUE text, opts;
     rb_scan_args(argc, argv, "1:", &text, &opts);
+    SE_CHECK_KEYWORDS(opts, id_max_tokens, id_validate_encoding);
 
     model_wrapper_t *w = get_model(self);
     Check_Type(text, T_STRING);
     check_text_encoding_mode(
         text, -1, resolve_encoding_validation(lookup_option(opts, id_validate_encoding)));
 
-    uint32_t max_tokens = resolve_max_tokens(&w->model, lookup_option(opts, id_max_tokens));
-
-    se_scratch_t scratch;
-    se_scratch_init(&scratch);
-    if (!se_scratch_reserve(&scratch, w->model.meta.dim)) {
-        se_scratch_free(&scratch);
+    tokenize_scratch_job_t job;
+    memset(&job, 0, sizeof(job));
+    job.model = &w->model;
+    job.input = (const uint8_t *)RSTRING_PTR(text);
+    job.input_len = (size_t)RSTRING_LEN(text);
+    job.max_tokens = resolve_max_tokens(&w->model, lookup_option(opts, id_max_tokens));
+    job.ids = Qnil;
+    se_error_clear(&job.err);
+
+    job.scratch = se_scratch_acquire(w->model.meta.dim);
+    if (!job.scratch)
         rb_raise(rb_eNoMemError, "out of memory");
-    }
 
-    se_token_stats_t stats;
-    se_error_t err;
-    se_error_clear(&err);
-    se_status_t rc = se_tokenize(&w->model, &scratch, (const uint8_t *)RSTRING_PTR(text),
-                                 (size_t)RSTRING_LEN(text), max_tokens, &stats, &err, NULL);
-    if (rc != SE_OK) {
-        se_scratch_free(&scratch);
-        raise_se(&err);
-    }
+    rb_ensure(tokenize_scratch_body, (VALUE)(uintptr_t)&job, tokenize_scratch_ensure,
+              (VALUE)(uintptr_t)&job);
+    RB_GC_GUARD(text);
 
-    VALUE ids = rb_ary_new_capa((long)stats.token_count);
-    for (uint32_t i = 0; i < stats.token_count; i++)
-        rb_ary_push(ids, UINT2NUM(scratch.ids[i]));
+    if (job.rc != SE_OK)
+        raise_se(&job.err);
 
-    se_scratch_free(&scratch);
-    RB_GC_GUARD(text);
-    return ids;
+    return job.ids;
 }
 
 typedef struct {
@@ -1028,17 +1137,15 @@ static VALUE num2ull_at_value(VALUE arg) {
 
 static void *ids_execute(void *arg) {
     ids_job_t *job = (ids_job_t *)arg;
-    se_scratch_t scratch;
-    se_scratch_init(&scratch);
-    if (!se_scratch_reserve(&scratch, job->model->meta.dim)) {
+    se_scratch_t *scratch = se_scratch_acquire(job->model->meta.dim);
+    if (!scratch) {
         se_error_set(&job->error, SE_ERR_OOM, "out of memory while sizing scratch buffers");
-        se_scratch_free(&scratch);
         return NULL;
     }
     se_error_clear(&job->error);
-    se_embed_ids(job->model, &scratch, job->ids, job->n_ids, job->out, &job->stats, &job->error,
+    se_embed_ids(job->model, scratch, job->ids, job->n_ids, job->out, &job->stats, &job->error,
                  &job->cancelled);
-    se_scratch_free(&scratch);
+    se_scratch_release(scratch);
     return NULL;
 }
 
@@ -1168,6 +1275,7 @@ static VALUE embed_token_ids_value(VALUE self, VALUE ids_value, VALUE max_tokens
 static VALUE model_embed_token_ids(int argc, VALUE *argv, VALUE self) {
     VALUE ids_value, opts;
     rb_scan_args(argc, argv, "1:", &ids_value, &opts);
+    SE_CHECK_KEYWORDS(opts, id_max_tokens, id_format, id_threads);
     reject_parallel_threads(opts);
     return embed_token_ids_value(self, ids_value, lookup_option(opts, id_max_tokens),
                                  resolve_vector_format(lookup_option(opts, id_format)), NULL);
@@ -1176,6 +1284,7 @@ static VALUE model_embed_token_ids(int argc, VALUE *argv, VALUE self) {
 static VALUE model_embed_token_ids_with_stats(int argc, VALUE *argv, VALUE self) {
     VALUE ids_value, opts;
     rb_scan_args(argc, argv, "1:", &ids_value, &opts);
+    SE_CHECK_KEYWORDS(opts, id_max_tokens, id_format, id_threads);
     reject_parallel_threads(opts);
 
     se_token_stats_t stats;
@@ -1315,6 +1424,7 @@ static void topk_check_matrix(VALUE matrix, size_t matrix_bytes, se_vector_forma
 static VALUE top_k_impl(int argc, VALUE *argv, VALUE self, int cosine) {
     VALUE query, matrix, k_val, opts;
     rb_scan_args(argc, argv, "3:", &query, &matrix, &k_val, &opts);
+    SE_CHECK_KEYWORDS(opts, id_dim, id_format, id_allow_unfrozen);
     (void)self;
 
     Check_Type(query, T_STRING);

data/lib/static_embeddings/version.rb

@@ -1,3 +1,3 @@
 module StaticEmbeddings
-  VERSION = "0.1.2"
+  VERSION = "0.1.4"
 end

data/static_embeddings.gemspec

@@ -25,8 +25,10 @@ Gem::Specification.new do |spec|
     "lib/models/*.semb",
     "exe/*",
     "ext/**/*.{c,h,rb}",
+    "benchmark/**/*.rb",
     "docs/**/*.md",
     "tools/*.rb",
+    "Rakefile",
     "static_embeddings.gemspec",
     "README.md",
     "CHANGELOG.md",

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: static_embeddings
 version: !ruby/object:Gem::Version
-  version: 0.1.2
+  version: 0.1.4
 platform: ruby
 authors:
 - Roman Haydarov
@@ -80,7 +80,13 @@ files:
 - CHANGELOG.md
 - LICENSE.txt
 - README.md
+- Rakefile
+- benchmark/_support.rb
+- benchmark/core_paths.rb
+- benchmark/gvl_threshold.rb
 - docs/ARCHITECTURE.md
+- docs/BENCHMARKING.md
+- docs/LIMITATIONS.md
 - docs/MODEL_AUDIT.md
 - docs/PERFORMANCE.md
 - exe/static_embeddings