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.



1151
1152
1153
1154
1155
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1151

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.



1089
1090
1091
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1089

def directory
  @directory
end

#readers ⇒ Object (readonly)

Returns the value of attribute readers.



1089
1090
1091
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1089

def readers
  @readers
end

#tensors ⇒ Object (readonly)

Returns the value of attribute tensors.



1089
1090
1091
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1089

def tensors
  @tensors
end

Class Method Details

.from_directory(directory) ⇒ Object

Raises:



1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1091

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 



1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1102

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 



1165
1166
1167
1168
1169
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1165

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

#fetch_any(names) ⇒ Object

Raises:



1171
1172
1173
1174
1175
1176
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1171

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) 


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

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

#names ⇒ Object 



1157
1158
1159
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 1157

def names
  tensors.keys
end