Class: Cohere::Transcribe::GGUF::Writer

Inherits:
Object
  • Object
show all
Defined in:
lib/cohere/transcribe/gguf_writer.rb

Overview

Produces a GGUF v3 document from metadata and streaming tensor writers.

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(alignment: DEFAULT_ALIGNMENT) ⇒ Writer

Returns a new instance of Writer.



50
51
52
53
54
55
56
57
58
59
# File 'lib/cohere/transcribe/gguf_writer.rb', line 50

def initialize(alignment: DEFAULT_ALIGNMENT)
  unless alignment.is_a?(Integer) && alignment.positive? && alignment.nobits?(alignment - 1)
    raise ArgumentError, "GGUF alignment must be a positive power of two"
  end

  @alignment = alignment
  @metadata = {}
  @tensors = []
  add_uint32("general.alignment", alignment) unless alignment == DEFAULT_ALIGNMENT
end

Instance Attribute Details

#alignmentObject (readonly)

Returns the value of attribute alignment.



48
49
50
# File 'lib/cohere/transcribe/gguf_writer.rb', line 48

def alignment
  @alignment
end

#metadataObject (readonly)

Returns the value of attribute metadata.



48
49
50
# File 'lib/cohere/transcribe/gguf_writer.rb', line 48

def 
  @metadata
end

#tensorsObject (readonly)

Returns the value of attribute tensors.



48
49
50
# File 'lib/cohere/transcribe/gguf_writer.rb', line 48

def tensors
  @tensors
end

Instance Method Details

#add_bool(key, value) ⇒ Object

Raises:

  • (ArgumentError)


77
78
79
80
81
# File 'lib/cohere/transcribe/gguf_writer.rb', line 77

def add_bool(key, value)
  raise ArgumentError, "GGUF bool metadata must be true or false" unless [true, false].include?(value)

  (key, :bool, value)
end

#add_float32(key, value) ⇒ Object



73
74
75
# File 'lib/cohere/transcribe/gguf_writer.rb', line 73

def add_float32(key, value)
  (key, :float32, Float(value))
end

#add_string(key, value) ⇒ Object



61
62
63
# File 'lib/cohere/transcribe/gguf_writer.rb', line 61

def add_string(key, value)
  (key, :string, String(value))
end

#add_string_array(key, values) ⇒ Object

Raises:

  • (ArgumentError)


83
84
85
86
87
88
# File 'lib/cohere/transcribe/gguf_writer.rb', line 83

def add_string_array(key, values)
  values = values.map { |value| String(value) }.freeze
  raise ArgumentError, "GGUF arrays must not be empty" if values.empty?

  (key, :array, values, element_type: :string)
end

#add_tensor(name, shape:, dtype:, &write_data) ⇒ Object

Raises:

  • (ArgumentError)


90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
# File 'lib/cohere/transcribe/gguf_writer.rb', line 90

def add_tensor(name, shape:, dtype:, &write_data)
  name = String(name)
  raise ArgumentError, "GGUF tensor name must not be empty" if name.empty?
  raise ArgumentError, "GGUF tensor #{name.inspect} has no data writer" unless write_data
  raise ArgumentError, "Duplicate GGUF tensor name #{name.inspect}" if tensors.any? { |tensor| tensor.name == name }

  shape = Array(shape)
  unless !shape.empty? && shape.all? { |dimension| dimension.is_a?(Integer) && dimension.positive? }
    raise ArgumentError, "GGUF tensor #{name.inspect} has an invalid shape"
  end

  dtype = dtype.to_sym
  width = TENSOR_WIDTHS[dtype]
  raise ArgumentError, "Unsupported GGUF tensor dtype #{dtype.inspect}" unless width

  elements = shape.reduce(1, :*)
  tensors << Tensor.new(
    name: name.freeze,
    shape: shape.freeze,
    dtype: dtype,
    nbytes: elements * width,
    write_data: write_data
  )
  self
end

#add_uint32(key, value) ⇒ Object



65
66
67
# File 'lib/cohere/transcribe/gguf_writer.rb', line 65

def add_uint32(key, value)
  (key, :uint32, unsigned_integer(value, bits: 32))
end

#add_uint64(key, value) ⇒ Object



69
70
71
# File 'lib/cohere/transcribe/gguf_writer.rb', line 69

def add_uint64(key, value)
  (key, :uint64, unsigned_integer(value, bits: 64))
end

#write(path, overwrite: false, fsync: true) ⇒ Object



116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
# File 'lib/cohere/transcribe/gguf_writer.rb', line 116

def write(path, overwrite: false, fsync: true)
  output = Pathname(path).expand_path
  raise Error, "Output directory #{output.dirname} does not exist" unless output.dirname.directory?
  raise Error, "Output file #{output} already exists" if output.exist? && !overwrite

  temporary = Tempfile.new([".#{output.basename}", ".tmp"], output.dirname.to_s)
  temporary.binmode
  begin
    write_to(temporary)
    temporary.flush
    temporary.fsync if fsync
    temporary.close
    raise Error, "Output file #{output} already exists" if output.exist? && !overwrite

    if overwrite
      File.rename(temporary.path, output)
    else
      # A same-directory hard link is an atomic no-replace publish on
      # POSIX filesystems. It closes the existence-check race where a
      # concurrent writer could otherwise be silently overwritten by
      # rename(2).
      begin
        File.link(temporary.path, output)
      rescue Errno::EEXIST
        raise Error, "Output file #{output} already exists"
      end
      temporary.unlink
    end
    sync_directory(output.dirname) if fsync
  ensure
    temporary.close unless temporary.closed?
    temporary.unlink
  end
  output
rescue Errno::EACCES, Errno::ENOSPC, Errno::EROFS => e
  raise Error, "Cannot write GGUF file #{output}: #{e.message}"
end

#write_to(io) ⇒ Object



154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
# File 'lib/cohere/transcribe/gguf_writer.rb', line 154

def write_to(io)
  validate_document!
  io.write(MAGIC)
  io.write([VERSION].pack("L<"))
  io.write([tensors.length, .length].pack("Q<Q<"))
  .each do |key, entry|
    write_string(io, key)
    io.write([VALUE_TYPES.fetch(entry.type)].pack("L<"))
    write_value(io, entry.type, entry.value, element_type: entry.element_type)
  end

  offset = 0
  tensors.each do |tensor|
    write_string(io, tensor.name)
    io.write([tensor.shape.length].pack("L<"))
    io.write(tensor.shape.reverse.pack("Q<*"))
    io.write([TENSOR_TYPES.fetch(tensor.dtype)].pack("L<"))
    io.write([offset].pack("Q<"))
    offset += padded(tensor.nbytes)
  end

  write_padding(io, io.pos)
  tensors.each do |tensor|
    before = io.pos
    tensor.write_data.call(io)
    actual = io.pos - before
    if actual != tensor.nbytes
      raise Error,
            "Tensor #{tensor.name.inspect} wrote #{actual} bytes; expected #{tensor.nbytes}"
    end
    write_padding(io, tensor.nbytes)
  end
  io
end