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.



1208
1209
1210
1211
1212
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1208

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.



1146
1147
1148
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1146

def directory
  @directory
end

#readers ⇒ Object (readonly)

Returns the value of attribute readers.



1146
1147
1148
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1146

def readers
  @readers
end

#tensors ⇒ Object (readonly)

Returns the value of attribute tensors.



1146
1147
1148
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1146

def tensors
  @tensors
end

Class Method Details

.from_directory(directory) ⇒ Object

Raises:



1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1148

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 



1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1159

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 



1222
1223
1224
1225
1226
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1222

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

#fetch_any(names) ⇒ Object

Raises:



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

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) 


1218
1219
1220
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1218

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

#names ⇒ Object 



1214
1215
1216
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1214

def names
  tensors.keys
end