Class: Cohere::Transcribe::PyTorchCheckpoint::TensorSet

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

Overview

One unsharded pytorch_model.bin or its Transformers shard index.

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.



1224
1225
1226
1227
1228
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1224

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.



1162
1163
1164
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1162

def directory
  @directory
end

#readers ⇒ Object (readonly)

Returns the value of attribute readers.



1162
1163
1164
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1162

def readers
  @readers
end

#tensors ⇒ Object (readonly)

Returns the value of attribute tensors.



1162
1163
1164
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1162

def tensors
  @tensors
end

Class Method Details

.from_directory(directory) ⇒ Object

Raises:



1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1164

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

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

  from_index(directory, index)
end

.from_index(directory, index_path) ⇒ Object 



1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1175

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, "PyTorch index #{index_path} has an invalid weight_map"
  end

  shard_names = weight_map.values.uniq
  shard_names.each { |name| validate_shard_name!(name, index_path) }
  readers = shard_names.to_h do |name|
    shard = directory.join(name)
    raise Error, "PyTorch shard #{shard} is missing" unless shard.file?

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

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

    raise Error,
          "PyTorch 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 PyTorch index #{index_path}: #{e.message}"
rescue Errno::EACCES, Errno::ENOENT => e
  raise Error, "Cannot read PyTorch index #{index_path}: #{e.message}"
end

Instance Method Details

#fetch(name) ⇒ Object 



1238
1239
1240
1241
1242
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1238

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

#fetch_any(names) ⇒ Object

Raises:



1244
1245
1246
1247
1248
1249
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1244

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) 


1234
1235
1236
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1234

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

#names ⇒ Object 



1230
1231
1232
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1230

def names
  tensors.keys
end