Module: StaticEmbeddings::Safetensors

Defined in:
lib/static_embeddings/safetensors.rb

Constant Summary collapse

MAX_HEADER_BYTES = 
100 * 1024 * 1024
SUPPORTED_DTYPES = 
{ "F32" => 4 }.freeze

Class Method Summary collapse

Class Method Details

.checked_offsets(name, offsets, body_size) ⇒ Object 



63
64
65
66
67
68
69
70
71
72
73
74
75
 
# File 'lib/static_embeddings/safetensors.rb', line 63

def checked_offsets(name, offsets, body_size)
  unless offsets.is_a?(Array) && offsets.length == 2
    raise ConversionError, "tensor #{name}: malformed data_offsets"
  end

  begin_off, end_off = offsets
  unless begin_off.is_a?(Integer) && end_off.is_a?(Integer) && begin_off >= 0 &&
         end_off >= begin_off && end_off <= body_size
    raise ConversionError, "tensor #{name}: data_offsets out of bounds"
  end

  [begin_off, end_off]
end

.checked_tensor_bytes(name, dtype, shape) ⇒ Object

Raises:



77
78
79
80
81
82
83
84
85
 
# File 'lib/static_embeddings/safetensors.rb', line 77

def checked_tensor_bytes(name, dtype, shape)
  element_size = SUPPORTED_DTYPES[dtype]
  raise ConversionError, "tensor #{name}: dtype #{dtype} is not supported (F32 only)" unless element_size
  unless shape.is_a?(Array) && shape.all? { |dim| dim.is_a?(Integer) && dim >= 0 }
    raise ConversionError, "tensor #{name}: malformed shape #{shape.inspect}"
  end

  shape.reduce(1, :*) * element_size
end

.parse_header(data) ⇒ Object

Raises:



31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
 
# File 'lib/static_embeddings/safetensors.rb', line 31

def parse_header(data)
  raise ConversionError, "safetensors file is shorter than its length prefix" if data.bytesize < 8

  header_len = data.byteslice(0, 8).unpack1("Q<")
  unless header_len.positive? && header_len <= MAX_HEADER_BYTES
    raise ConversionError, "implausible safetensors header length #{header_len}"
  end

  body_offset = 8 + header_len
  raise ConversionError, "safetensors header runs past end of file" if body_offset > data.bytesize

  raw_header = data.byteslice(8, header_len)
  raise ConversionError, "safetensors header is not a JSON object" unless raw_header.lstrip.start_with?("{")

  [JSON.parse(raw_header), body_offset]
end

.read(path) ⇒ Object 



10
11
12
13
14
15
16
17
18
19
20
21
22
 
# File 'lib/static_embeddings/safetensors.rb', line 10

def read(path)
  data = File.binread(path)
  header, body_offset = parse_header(data)
  body_size = data.bytesize - body_offset

  tensors = header.each_with_object({}) do |(name, spec), acc|
    next if name == "__metadata__"

    acc[name] = tensor_from(data, body_offset, body_size, name, spec)
  end

  { metadata: header["__metadata__"] || {}, tensors: tensors }
end

.tensor_from(data, body_offset, body_size, name, spec) ⇒ Object 



48
49
50
51
52
53
54
55
56
57
58
59
60
61
 
# File 'lib/static_embeddings/safetensors.rb', line 48

def tensor_from(data, body_offset, body_size, name, spec)
  dtype = spec["dtype"]
  shape = spec["shape"]
  begin_off, end_off = checked_offsets(name, spec["data_offsets"], body_size)
  expected = checked_tensor_bytes(name, dtype, shape)
  actual = end_off - begin_off

  if actual != expected
    raise ConversionError,
          "tensor #{name}: shape #{shape.inspect} implies #{expected} bytes, offsets span #{actual}"
  end

  { dtype: dtype, shape: shape, bytes: data.byteslice(body_offset + begin_off, actual) }
end

.write(path, name, shape, floats) ⇒ Object 



24
25
26
27
28
29
 
# File 'lib/static_embeddings/safetensors.rb', line 24

def write(path, name, shape, floats)
  body = floats.pack("e*")
  header = JSON.generate(name => { "dtype" => "F32", "shape" => shape, "data_offsets" => [0, body.bytesize] })
  padded = header << (" " * ((8 - (header.bytesize % 8)) % 8))
  File.binwrite(path, [padded.bytesize].pack("Q<") << padded << body)
end