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 



402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
 
# File 'ext/static_embeddings/static_embeddings.c', line 402

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 



420
421
422
423
424
425
426
427
428
 
# File 'ext/static_embeddings/static_embeddings.c', line 420

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) 


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

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 



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

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 



1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
 
# File 'ext/static_embeddings/static_embeddings.c', line 1006

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



916
917
918
919
920
921
922
923
924
925
926
927
 
# File 'ext/static_embeddings/static_embeddings.c', line 916

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



1275
1276
1277
1278
1279
1280
1281
1282
 
# File 'ext/static_embeddings/static_embeddings.c', line 1275

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

#embed_token_ids_with_stats(*args) ⇒ Object 



1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
 
# File 'ext/static_embeddings/static_embeddings.c', line 1284

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



1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
 
# File 'ext/static_embeddings/static_embeddings.c', line 1017

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


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

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

#mapped_bytes ⇒ Object 



467
468
469
 
# File 'ext/static_embeddings/static_embeddings.c', line 467

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

#max_tokens ⇒ Object 



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

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) 


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

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 



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

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 



1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
 
# File 'ext/static_embeddings/static_embeddings.c', line 1069

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

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

    rb_ensure(tokenize_scratch_body, (VALUE)(uintptr_t)&job, tokenize_scratch_ensure,
              (VALUE)(uintptr_t)&job);
    RB_GC_GUARD(text);

    if (job.rc != SE_OK)
        raise_se(&job.err);

    return job.ids;
}

#unk_id ⇒ Object 



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

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

#vocab_size ⇒ Object 



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

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

#warmup! ⇒ Object 



482
483
484
485
486
487
488
489
490
 
# File 'ext/static_embeddings/static_embeddings.c', line 482

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