Class: Cohere::Transcribe::Safetensors::TensorSet

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

Overview

Presents one unsharded file or a model.safetensors.index.json shard set through the same tensor-name lookup API.

Instance Attribute Summary collapse

Class Method Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(directory, readers, tensors: nil) ⇒ TensorSet

Returns a new instance of TensorSet.



255
256
257
258
259
 
# File 'lib/cohere/transcribe/safetensors.rb', line 255

def initialize(directory, readers, tensors: nil)
  @directory = Pathname(directory).expand_path
  @readers = readers.freeze
  @tensors = (tensors || readers.flat_map { |reader| reader.tensors.to_a }.to_h).freeze
end

Instance Attribute Details

#directory ⇒ Object (readonly)

Returns the value of attribute directory.



198
199
200
 
# File 'lib/cohere/transcribe/safetensors.rb', line 198

def directory
  @directory
end

#readers ⇒ Object (readonly)

Returns the value of attribute readers.



198
199
200
 
# File 'lib/cohere/transcribe/safetensors.rb', line 198

def readers
  @readers
end

#tensors ⇒ Object (readonly)

Returns the value of attribute tensors.



198
199
200
 
# File 'lib/cohere/transcribe/safetensors.rb', line 198

def tensors
  @tensors
end

Class Method Details

.from_directory(directory) ⇒ Object

Raises:



200
201
202
203
204
205
206
207
208
209
 
# File 'lib/cohere/transcribe/safetensors.rb', line 200

def self.from_directory(directory)
  directory = Pathname(directory).expand_path
  single = directory.join("model.safetensors")
  return new(directory, [Reader.new(single)]) if single.file?

  index = directory.join("model.safetensors.index.json")
  raise Error, "#{directory} does not contain model.safetensors or its index" unless index.file?

  from_index(directory, index)
end

.from_index(directory, index_path) ⇒ Object 



211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
 
# File 'lib/cohere/transcribe/safetensors.rb', line 211

def self.from_index(directory, index_path)
  payload = JSON.parse(index_path.read(encoding: "UTF-8"))
  weight_map = payload["weight_map"]
  unless weight_map.is_a?(Hash) && !weight_map.empty? &&
         weight_map.all? { |name, file| name.is_a?(String) && file.is_a?(String) }
    raise Error, "Safetensors index #{index_path} has an invalid weight_map"
  end

  shard_names = weight_map.values.uniq
  shard_names.each do |name|
    candidate = Pathname(name)
    if name.empty? || candidate.absolute? || candidate.each_filename.any? { |part| part == ".." }
      raise Error, "Safetensors index #{index_path} contains an invalid shard path #{name.inspect}"
    end
  end
  readers = shard_names.to_h do |name|
    path = directory.join(name)
    raise Error, "Safetensors shard #{path} is missing" unless path.file?

    [name, Reader.new(path)]
  end

  mapped = weight_map.to_h do |name, shard|
    reader = readers.fetch(shard)
    raise Error, "Safetensors index maps missing tensor #{name.inspect} to #{shard}" unless reader.key?(name)

    [name, reader.fetch(name)]
  end
  readers.each do |shard, reader|
    indexed_names = weight_map.filter_map { |name, filename| name if filename == shard }
    unindexed = reader.names - indexed_names
    next if unindexed.empty?

    raise Error,
          "Safetensors shard #{reader.path} contains tensors absent from its index: " \
          "#{unindexed.first(8).map(&:inspect).join(", ")}"
  end
  new(directory, readers.values, tensors: mapped)
rescue JSON::ParserError, EncodingError => e
  raise Error, "Invalid safetensors index #{index_path}: #{e.message}"
rescue Errno::EACCES, Errno::ENOENT => e
  raise Error, "Cannot read safetensors index #{index_path}: #{e.message}"
end

Instance Method Details

#fetch(name) ⇒ Object 



269
270
271
272
273
 
# File 'lib/cohere/transcribe/safetensors.rb', line 269

def fetch(name)
  tensors.fetch(name)
rescue KeyError
  raise Error, "Tensor #{name.inspect} is not present in #{directory}"
end

#fetch_any(names) ⇒ Object

Raises:



275
276
277
278
279
280
 
# File 'lib/cohere/transcribe/safetensors.rb', line 275

def fetch_any(names)
  name = names.find { |candidate| tensors.key?(candidate) }
  return tensors.fetch(name) if name

  raise Error, "None of the tensor aliases are present: #{names.map(&:inspect).join(", ")}"
end

#key?(name) ⇒ Boolean

Returns:

  • (Boolean) 


265
266
267
 
# File 'lib/cohere/transcribe/safetensors.rb', line 265

def key?(name)
  tensors.key?(name)
end

#names ⇒ Object 



261
262
263
 
# File 'lib/cohere/transcribe/safetensors.rb', line 261

def names
  tensors.keys
end