Module: StaticEmbeddings

Defined in:
lib/static_embeddings.rb,
lib/static_embeddings/cli.rb,
lib/static_embeddings/model.rb,
lib/static_embeddings/paths.rb,
lib/static_embeddings/errors.rb,
lib/static_embeddings/format.rb,
lib/static_embeddings/version.rb,
lib/static_embeddings/converter.rb,
lib/static_embeddings/reference.rb,
lib/static_embeddings/safetensors.rb,
lib/static_embeddings/unicode_tables.rb,
ext/static_embeddings/static_embeddings.c

Defined Under Namespace

Modules: Format, Paths, Safetensors, UnicodeTables Classes: CLI, ConversionError, Converter, EmptyInputError, EncodingError, Error, InvalidModelError, Model, ModelNotFound, Reference, UnsupportedModelError

Constant Summary collapse

VERSION = 
"0.1.2"
FORMAT_VERSION = 
UINT2NUM(SE_FORMAT_VERSION)
TOKENIZER_BERT_WORDPIECE_V1 = 
UINT2NUM(SE_TOKENIZER_BERT_WORDPIECE_V1)

Class Method Summary collapse

Class Method Details

.__alloc_stats__ ⇒ Object 



1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
 
# File 'ext/static_embeddings/static_embeddings.c', line 1450

static VALUE se_alloc_stats_hash(VALUE self) {
    (void)self;
    se_alloc_stats_t stats[SE_ALLOC_CATEGORY_COUNT];
    se_alloc_stats_snapshot(stats);

    VALUE out = rb_hash_new();
    ID id_current_bytes = rb_intern("current_bytes");
    ID id_peak_bytes = rb_intern("peak_bytes");
    ID id_total_allocated_bytes = rb_intern("total_allocated_bytes");
    ID id_total_freed_bytes = rb_intern("total_freed_bytes");
    ID id_alloc_count = rb_intern("alloc_count");
    ID id_realloc_count = rb_intern("realloc_count");
    ID id_free_count = rb_intern("free_count");

    for (int i = 0; i < SE_ALLOC_CATEGORY_COUNT; i++) {
        VALUE item = rb_hash_new();
        rb_hash_aset(item, ID2SYM(id_current_bytes), SIZET2NUM(stats[i].current_bytes));
        rb_hash_aset(item, ID2SYM(id_peak_bytes), SIZET2NUM(stats[i].peak_bytes));
        rb_hash_aset(item, ID2SYM(id_total_allocated_bytes),
                     SIZET2NUM(stats[i].total_allocated_bytes));
        rb_hash_aset(item, ID2SYM(id_total_freed_bytes), SIZET2NUM(stats[i].total_freed_bytes));
        rb_hash_aset(item, ID2SYM(id_alloc_count), SIZET2NUM(stats[i].alloc_count));
        rb_hash_aset(item, ID2SYM(id_realloc_count), SIZET2NUM(stats[i].realloc_count));
        rb_hash_aset(item, ID2SYM(id_free_count), SIZET2NUM(stats[i].free_count));
        rb_hash_aset(out, ID2SYM(rb_intern(se_alloc_category_name((se_alloc_category_t)i))), item);
    }

    return out;
}

.__alloc_stats_reset__ ⇒ Object 



1480
1481
1482
1483
1484
 
# File 'ext/static_embeddings/static_embeddings.c', line 1480

static VALUE se_alloc_stats_reset_bang(VALUE self) {
    (void)self;
    se_alloc_stats_reset();
    return Qnil;
}

.builtin_available?(name = :demo) ⇒ Boolean

Returns:

  • (Boolean) 


45
46
47
 
# File 'lib/static_embeddings.rb', line 45

def builtin_available?(name = :demo)
  Paths.builtin_available?(name)
end

.cache_dir ⇒ Object 



49
50
51
 
# File 'lib/static_embeddings.rb', line 49

def cache_dir
  Paths.cache_dir
end

.convert(source_dir, output_path:, model_id: nil, max_tokens: Converter::REFERENCE_MAX_TOKENS) ⇒ Object 



67
68
69
70
71
 
# File 'lib/static_embeddings.rb', line 67

def convert(source_dir, output_path:, model_id: nil, max_tokens: Converter::REFERENCE_MAX_TOKENS)
  converter = Converter.new(source_dir)
  converter.convert(output_path: output_path, model_id: model_id, max_tokens: max_tokens)
  converter.report
end

.cosine_top_k(*args) ⇒ Object 



1406
1407
1408
 
# File 'ext/static_embeddings/static_embeddings.c', line 1406

static VALUE se_cosine_top_k(int argc, VALUE *argv, VALUE self) {
    return top_k_impl(argc, argv, self, 1);
}

.decode_f16(blob) ⇒ Object 



1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
 
# File 'ext/static_embeddings/static_embeddings.c', line 1432

static VALUE se_decode_f16(VALUE self, VALUE blob) {
    (void)self;
    Check_Type(blob, T_STRING);

    long n = RSTRING_LEN(blob);
    if (n % 2 != 0)
        rb_raise(rb_eArgError, "f16 blob byte size must be a multiple of 2");

    VALUE out = rb_ary_new_capa(n / 2);
    for (long i = 0; i < n / 2; i++) {
        const uint8_t *src = (const uint8_t *)RSTRING_PTR(blob) + (size_t)i * 2;
        rb_ary_push(out, DBL2NUM((double)se_read_f16le(src)));
    }
    RB_GC_GUARD(blob);
    return out;
}

.dot_top_k(*args) ⇒ Object 



1410
1411
1412
 
# File 'ext/static_embeddings/static_embeddings.c', line 1410

static VALUE se_dot_top_k(int argc, VALUE *argv, VALUE self) {
    return top_k_impl(argc, argv, self, 0);
}

.encode_f16(ary) ⇒ Object 



1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
 
# File 'ext/static_embeddings/static_embeddings.c', line 1414

static VALUE se_encode_f16(VALUE self, VALUE ary) {
    (void)self;
    Check_Type(ary, T_ARRAY);

    long n = RARRAY_LEN(ary);
    size_t bytes;
    if (!se_checked_mul_size((size_t)n, 2, &bytes) || !se_size_fits_long(bytes))
        rb_raise(rb_eArgError, "vector is too large");

    VALUE out = rb_str_new(NULL, (long)bytes);
    rb_enc_associate(out, binary_encoding);
    for (long i = 0; i < n; i++) {
        double v = NUM2DBL(rb_ary_entry(ary, i));
        se_write_f16le((uint8_t *)RSTRING_PTR(out) + (size_t)i * 2, (float)v);
    }
    return out;
}

.load(path, verify: false) ⇒ Object

Raises:



29
30
31
32
33
34
35
36
37
38
39
 
# File 'lib/static_embeddings.rb', line 29

def load(path, verify: false)
  expanded = File.expand_path(path.to_s)
  raise ModelNotFound, "no model at #{expanded}" unless File.file?(expanded)

  if verify
    result = Format.verify(expanded)
    raise InvalidModelError, "checksum mismatch for #{expanded}" unless result[:ok]
  end

  Model.new(expanded)
end

.load_builtin(name = :demo, verify: false) ⇒ Object 



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

def load_builtin(name = :demo, verify: false)
  load(Paths.builtin_path(name), verify: verify)
end

.load_model(model_id, verify: false) ⇒ Object 



57
58
59
60
61
62
63
64
65
 
# File 'lib/static_embeddings.rb', line 57

def load_model(model_id, verify: false)
  path = model_path(model_id)
  unless File.file?(path)
    raise ModelNotFound,
          "model #{model_id.inspect} is not installed. " \
          "Convert it first: static_embeddings convert <hf-dir> --id #{model_id}"
  end
  load(path, verify: verify)
end

.model_path(model_id) ⇒ Object 



53
54
55
 
# File 'lib/static_embeddings.rb', line 53

def model_path(model_id)
  Paths.model_path(model_id)
end

.normalize_format(format) ⇒ Object 



106
107
108
109
110
111
112
113
114
115
 
# File 'lib/static_embeddings.rb', line 106

def normalize_format(format)
  case format&.to_sym
  when nil, :f32, :float32
    :f32
  when :f16, :float16
    :f16
  else
    raise ArgumentError, "unsupported embedding format #{format.inspect} (expected :f32 or :f16)"
  end
end

.pack(rows, format: :f32) ⇒ Object 



97
98
99
100
101
102
103
104
 
# File 'lib/static_embeddings.rb', line 97

def pack(rows, format: :f32)
  flat = rows.first.is_a?(Array) ? rows.flatten(1) : rows

  case normalize_format(format)
  when :f32 then flat.map(&:to_f).pack("e*")
  when :f16 then encode_f16(flat.map(&:to_f))
  end
end

.simd_backend ⇒ Object 



150
151
152
153
154
155
156
157
158
159
160
 
# File 'ext/static_embeddings/static_embeddings.c', line 150

static VALUE se_simd_backend(VALUE self) {
    (void)self;
    switch (se_current_f16_backend()) {
    case SE_F16_BACKEND_NEON_FP16:
        return rb_str_new_cstr("neon-fp16");
    case SE_F16_BACKEND_F16C:
        return rb_str_new_cstr("f16c");
    default:
        return rb_str_new_cstr("lut");
    }
}

.unpack(blob, dim, format: :f32) ⇒ Object

Raises:

  • (ArgumentError) 


77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
 
# File 'lib/static_embeddings.rb', line 77

def unpack(blob, dim, format: :f32)
  raise ArgumentError, "dim must be positive" unless dim.to_i.positive?

  floats =
    case normalize_format(format)
    when :f32
      raise ArgumentError, "f32 blob byte size must be a multiple of 4" unless (blob.bytesize % 4).zero?

      blob.unpack("e*")
    when :f16
      raise ArgumentError, "f16 blob byte size must be a multiple of 2" unless (blob.bytesize % 2).zero?

      decode_f16(blob)
    end

  raise ArgumentError, "blob is not a multiple of dim" unless (floats.length % dim).zero?

  floats.each_slice(dim).to_a
end

.verify(path) ⇒ Object 



73
74
75
 
# File 'lib/static_embeddings.rb', line 73

def verify(path)
  Format.verify(File.expand_path(path.to_s))
end