Class: StaticEmbeddings::Model

Inherits:
Object
  • Object
show all
Defined in:
lib/static_embeddings/model.rb,
ext/static_embeddings/static_embeddings.c

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(path) ⇒ Object 



398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
 
# File 'ext/static_embeddings/static_embeddings.c', line 398

static VALUE model_initialize(VALUE self, VALUE path) {
    model_wrapper_t *w;
    TypedData_Get_Struct(self, model_wrapper_t, &model_type, w);

    Check_Type(path, T_STRING);
    if (memchr(RSTRING_PTR(path), '\0', (size_t)RSTRING_LEN(path)))
        rb_raise(rb_eArgError, "path contains a null byte");

    se_error_t err;
    se_error_clear(&err);
    if (se_model_open(&w->model, StringValueCStr(path), &err) != SE_OK)
        raise_se(&err);

    w->open = 1;
    rb_ivar_set(self, rb_intern("@path"), rb_str_new_frozen(path));
    return self;
}

Instance Attribute Details

#path ⇒ Object (readonly)

Returns the value of attribute path.



5
6
7
 
# File 'lib/static_embeddings/model.rb', line 5

def path
  @path
end

Instance Method Details

#close ⇒ Object 



416
417
418
419
420
421
422
423
424
 
# File 'ext/static_embeddings/static_embeddings.c', line 416

static VALUE model_close(VALUE self) {
    model_wrapper_t *w;
    TypedData_Get_Struct(self, model_wrapper_t, &model_type, w);
    if (w->open) {
        se_model_close(&w->model);
        w->open = 0;
    }
    return Qnil;
}

#closed? ⇒ Boolean

Returns:

  • (Boolean) 


426
427
428
429
430
 
# File 'ext/static_embeddings/static_embeddings.c', line 426

static VALUE model_closed_p(VALUE self) {
    model_wrapper_t *w;
    TypedData_Get_Struct(self, model_wrapper_t, &model_type, w);
    return w->open ? Qfalse : Qtrue;
}

#cosine_top_k(query_blob, matrix_blob, k, **opts) ⇒ Object

Raises:

  • (ArgumentError) 


30
31
32
33
34
 
# File 'lib/static_embeddings/model.rb', line 30

def cosine_top_k(query_blob, matrix_blob, k, **opts)
  raise ArgumentError, "dim: is set by the model" if opts.key?(:dim)

  StaticEmbeddings.cosine_top_k(query_blob, matrix_blob, k, **opts.merge(dim: dim))
end

#dim ⇒ Object 



432
433
434
 
# File 'ext/static_embeddings/static_embeddings.c', line 432

static VALUE model_dim(VALUE self) {
    return UINT2NUM(get_model(self)->model.meta.dim);
}

#dot_top_k(query_blob, matrix_blob, k, **opts) ⇒ Object

Raises:

  • (ArgumentError) 


36
37
38
39
40
 
# File 'lib/static_embeddings/model.rb', line 36

def dot_top_k(query_blob, matrix_blob, k, **opts)
  raise ArgumentError, "dim: is set by the model" if opts.key?(:dim)

  StaticEmbeddings.dot_top_k(query_blob, matrix_blob, k, **opts.merge(dim: dim))
end

#embed(*args) ⇒ Object 



927
928
929
930
931
932
933
934
935
 
# File 'ext/static_embeddings/static_embeddings.c', line 927

static VALUE model_embed(int argc, VALUE *argv, VALUE self) {
    VALUE text, opts;
    rb_scan_args(argc, argv, "1:", &text, &opts);
    reject_parallel_threads(opts);
    return embed_one_value(self, text, lookup_option(opts, id_max_tokens),
                           resolve_vector_format(lookup_option(opts, id_format)),
                           resolve_encoding_validation(lookup_option(opts, id_validate_encoding)),
                           NULL);
}

#embed_array(text, **opts) ⇒ Object 



20
21
22
23
 
# File 'lib/static_embeddings/model.rb', line 20

def embed_array(text, **opts)
  format = opts.key?(:format) ? opts[:format] : :f32
  StaticEmbeddings.unpack(embed(text, **opts), dim, format: format).first
end

#embed_batch(*args) ⇒ Object 



867
868
869
870
871
872
873
874
875
876
877
 
# File 'ext/static_embeddings/static_embeddings.c', line 867

static VALUE model_embed_batch(int argc, VALUE *argv, VALUE self) {
    VALUE texts, opts;
    rb_scan_args(argc, argv, "1:", &texts, &opts);
    reject_parallel_threads(opts);

    VALUE max_tokens = lookup_option(opts, id_max_tokens);
    se_vector_format_t format = resolve_vector_format(lookup_option(opts, id_format));
    se_encoding_validation_t validation =
        resolve_encoding_validation(lookup_option(opts, id_validate_encoding));
    return embed_batch_internal(self, texts, max_tokens, format, validation);
}

#embed_batch_arrays(texts, **opts) ⇒ Object 



25
26
27
28
 
# File 'lib/static_embeddings/model.rb', line 25

def embed_batch_arrays(texts, **opts)
  format = opts.key?(:format) ? opts[:format] : :f32
  StaticEmbeddings.unpack(embed_batch(texts, **opts), dim, format: format)
end

#embed_token_ids(*args) ⇒ Object 



1168
1169
1170
1171
1172
1173
1174
 
# File 'ext/static_embeddings/static_embeddings.c', line 1168

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);
    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);
}

#embed_token_ids_with_stats(*args) ⇒ Object 



1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
 
# File 'ext/static_embeddings/static_embeddings.c', line 1176

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);
    reject_parallel_threads(opts);

    se_token_stats_t stats;
    memset(&stats, 0, sizeof(stats));
    VALUE vector =
        embed_token_ids_value(self, ids_value, lookup_option(opts, id_max_tokens),
                              resolve_vector_format(lookup_option(opts, id_format)), &stats);

    VALUE hash = rb_hash_new();
    rb_hash_aset(hash, ID2SYM(id_vector), vector);
    rb_hash_aset(hash, ID2SYM(id_token_count), UINT2NUM(stats.token_count));
    rb_hash_aset(hash, ID2SYM(id_unk_count), UINT2NUM(stats.unk_count));
    rb_hash_aset(hash, ID2SYM(id_truncated), stats.truncated ? Qtrue : Qfalse);
    return hash;
}

#embed_with_stats(*args) ⇒ Object 



937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
 
# File 'ext/static_embeddings/static_embeddings.c', line 937

static VALUE model_embed_with_stats(int argc, VALUE *argv, VALUE self) {
    VALUE text, opts;
    rb_scan_args(argc, argv, "1:", &text, &opts);
    reject_parallel_threads(opts);

    se_token_stats_t stats;
    VALUE vector = embed_one_value(
        self, text, lookup_option(opts, id_max_tokens),
        resolve_vector_format(lookup_option(opts, id_format)),
        resolve_encoding_validation(lookup_option(opts, id_validate_encoding)), &stats);

    VALUE hash = rb_hash_new();
    rb_hash_aset(hash, ID2SYM(id_vector), vector);
    rb_hash_aset(hash, ID2SYM(id_token_count), UINT2NUM(stats.token_count));
    rb_hash_aset(hash, ID2SYM(id_unk_count), UINT2NUM(stats.unk_count));
    rb_hash_aset(hash, ID2SYM(id_truncated), stats.truncated ? Qtrue : Qfalse);
    return hash;
}

#lowercase? ⇒ Boolean

Returns:

  • (Boolean) 


448
449
450
 
# File 'ext/static_embeddings/static_embeddings.c', line 448

static VALUE model_lowercase_p(VALUE self) {
    return get_model(self)->model.meta.do_lower_case ? Qtrue : Qfalse;
}

#mapped_bytes ⇒ Object 



463
464
465
 
# File 'ext/static_embeddings/static_embeddings.c', line 463

static VALUE model_mapped_bytes(VALUE self) {
    return SIZET2NUM(get_model(self)->model.map_size);
}

#max_tokens ⇒ Object 



440
441
442
 
# File 'ext/static_embeddings/static_embeddings.c', line 440

static VALUE model_max_tokens(VALUE self) {
    return UINT2NUM(get_model(self)->model.meta.max_tokens_default);
}

#model_id ⇒ Object 



16
17
18
 
# File 'lib/static_embeddings/model.rb', line 16

def model_id
  provenance["source_model_id"]
end

#normalized? ⇒ Boolean

Returns:

  • (Boolean) 


444
445
446
 
# File 'ext/static_embeddings/static_embeddings.c', line 444

static VALUE model_normalized_p(VALUE self) {
    return get_model(self)->model.meta.normalization_type == SE_NORMALIZATION_L2 ? Qtrue : Qfalse;
}

#provenance ⇒ Object 



7
8
9
10
11
12
13
14
 
# File 'lib/static_embeddings/model.rb', line 7

def provenance
  raw = provenance_json
  return {} if raw.nil?

  JSON.parse(raw)
rescue JSON::ParserError, EncodingError => e
  raise InvalidModelError, "invalid provenance JSON: #{e.message}"
end

#provenance_json ⇒ Object 



456
457
458
459
460
461
 
# File 'ext/static_embeddings/static_embeddings.c', line 456

static VALUE model_provenance_json(VALUE self) {
    model_wrapper_t *w = get_model(self);
    if (!w->model.provenance)
        return Qnil;
    return rb_enc_str_new(w->model.provenance, (long)w->model.provenance_size, utf8_encoding);
}

#to_s ⇒ Object Also known as: inspect 



42
43
44
 
# File 'lib/static_embeddings/model.rb', line 42

def to_s
  "#<StaticEmbeddings::Model #{model_id || path} dim=#{dim} vocab=#{vocab_size}>"
end

#tokenize(*args) ⇒ Object 



956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
 
# File 'ext/static_embeddings/static_embeddings.c', line 956

static VALUE model_tokenize(int argc, VALUE *argv, VALUE self) {
    VALUE text, opts;
    rb_scan_args(argc, argv, "1:", &text, &opts);

    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);
        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);
    }

    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]));

    se_scratch_free(&scratch);
    RB_GC_GUARD(text);
    return ids;
}

#unk_id ⇒ Object 



452
453
454
 
# File 'ext/static_embeddings/static_embeddings.c', line 452

static VALUE model_unk_id(VALUE self) {
    return UINT2NUM(get_model(self)->model.meta.unk_id);
}

#vocab_size ⇒ Object 



436
437
438
 
# File 'ext/static_embeddings/static_embeddings.c', line 436

static VALUE model_vocab_size(VALUE self) {
    return UINT2NUM(get_model(self)->model.meta.vocab_size);
}

#warmup! ⇒ Object 



478
479
480
481
482
483
484
485
486
 
# File 'ext/static_embeddings/static_embeddings.c', line 478

static VALUE model_warmup(VALUE self) {
    model_wrapper_t *w = get_model(self);
    warmup_job_t job;
    job.model = &w->model;
    job.pages = 0;
    rb_thread_call_without_gvl(warmup_execute, &job, NULL, NULL);
    RB_GC_GUARD(self);
    return self;
}