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 



478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
 
# File 'ext/static_embeddings/static_embeddings.c', line 478

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 



496
497
498
499
500
501
502
503
504
 
# File 'ext/static_embeddings/static_embeddings.c', line 496

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) 


506
507
508
509
510
 
# File 'ext/static_embeddings/static_embeddings.c', line 506

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

#dim ⇒ Object 



512
513
514
 
# File 'ext/static_embeddings/static_embeddings.c', line 512

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

#embed(*args) ⇒ Object 



986
987
988
989
990
991
992
 
# File 'ext/static_embeddings/static_embeddings.c', line 986

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)), 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 



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

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));
    return embed_batch_internal(self, texts, max_tokens, format);
}

#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 



1217
1218
1219
1220
1221
1222
1223
 
# File 'ext/static_embeddings/static_embeddings.c', line 1217

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 



1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
 
# File 'ext/static_embeddings/static_embeddings.c', line 1225

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 



994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
 
# File 'ext/static_embeddings/static_embeddings.c', line 994

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)), &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) 


528
529
530
 
# File 'ext/static_embeddings/static_embeddings.c', line 528

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

#mapped_bytes ⇒ Object 



543
544
545
 
# File 'ext/static_embeddings/static_embeddings.c', line 543

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

#max_tokens ⇒ Object 



520
521
522
 
# File 'ext/static_embeddings/static_embeddings.c', line 520

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) 


524
525
526
 
# File 'ext/static_embeddings/static_embeddings.c', line 524

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 



536
537
538
539
540
541
 
# File 'ext/static_embeddings/static_embeddings.c', line 536

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 



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

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

#tokenize(*args) ⇒ Object 



1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
 
# File 'ext/static_embeddings/static_embeddings.c', line 1011

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(text);

    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, (size_t)RSTRING_LEN(text), 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 



532
533
534
 
# File 'ext/static_embeddings/static_embeddings.c', line 532

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

#vocab_size ⇒ Object 



516
517
518
 
# File 'ext/static_embeddings/static_embeddings.c', line 516

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

#warmup! ⇒ Object 



558
559
560
561
562
563
564
565
566
 
# File 'ext/static_embeddings/static_embeddings.c', line 558

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