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.0"
FORMAT_VERSION = 
UINT2NUM(SE_FORMAT_VERSION)
TOKENIZER_BERT_WORDPIECE_V1 = 
UINT2NUM(SE_TOKENIZER_BERT_WORDPIECE_V1)

Class Method Summary collapse

Class Method Details

.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 



1659
1660
1661
 
# File 'ext/static_embeddings/static_embeddings.c', line 1659

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

.dot_top_k(*args) ⇒ Object 



1663
1664
1665
 
# File 'ext/static_embeddings/static_embeddings.c', line 1663

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

.half_to_float(half) ⇒ Object 



108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
 
# File 'lib/static_embeddings.rb', line 108

def half_to_float(half)
  sign = (half & 0x8000) << 16
  exp = (half >> 10) & 0x1f
  mant = half & 0x03ff

  bits = if exp.zero?
    if mant.zero?
      sign
    else
      exp = 1
      while (mant & 0x0400).zero?
        mant <<= 1
        exp -= 1
      end
      mant &= 0x03ff
      sign | ((exp + (127 - 15)) << 23) | (mant << 13)
    end
  elsif exp == 31
    sign | 0x7f800000 | (mant << 13)
  else
    sign | ((exp + (127 - 15)) << 23) | (mant << 13)
  end

  [bits].pack("L<").unpack1("e")
end

.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 



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

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

.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?

      blob.unpack("S<*").map { |half| half_to_float(half) }
    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