Class: Cohere::Transcribe::PyTorchCheckpoint::Reader

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

Overview

One checkpoint shard. Tensor bytes are never materialized as a whole.

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(path) ⇒ Reader

Returns a new instance of Reader.



777
778
779
780
781
782
783
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 777

def initialize(path)
  @path = Pathname(path).expand_path
  @storages = {}
  @storage_payloads = {}
  @storage_verification_mutex = Mutex.new
  @tensors = read_checkpoint.freeze
end

Instance Attribute Details

#path ⇒ Object (readonly)

Returns the value of attribute path.



775
776
777
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 775

def path
  @path
end

#tensors ⇒ Object (readonly)

Returns the value of attribute tensors.



775
776
777
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 775

def tensors
  @tensors
end

Instance Method Details

#fetch(name) ⇒ Object 



793
794
795
796
797
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 793

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

#key?(name) ⇒ Boolean

Returns:

  • (Boolean) 


789
790
791
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 789

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

#names ⇒ Object 



785
786
787
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 785

def names
  tensors.keys
end

#write_tensor(tensor, output, target_dtype:, converter: Safetensors::DTypeConverter.default, chunk_bytes: Safetensors::DEFAULT_CHUNK_BYTES) ⇒ Object

Raises:

  • (ArgumentError) 


799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
 
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 799

def write_tensor(tensor, output, target_dtype:, converter: Safetensors::DTypeConverter.default,
                 chunk_bytes: Safetensors::DEFAULT_CHUNK_BYTES)
  raise ArgumentError, "Tensor #{tensor.name.inspect} belongs to a different PyTorch reader" unless tensor.reader.equal?(self)

  target_dtype = target_dtype.to_s.upcase
  unless FLOAT_DTYPES.include?(tensor.dtype) && FLOAT_DTYPES.include?(target_dtype)
    raise Error, "Cannot convert #{tensor.name.inspect} from #{tensor.dtype} to #{target_dtype}"
  end
  raise ArgumentError, "chunk_bytes must be positive" unless chunk_bytes.positive?

  source_width = Safetensors::DTYPE_BYTES.fetch(tensor.dtype)
  elements_per_chunk = [chunk_bytes / source_width, 1].max
  bytes_per_chunk = elements_per_chunk * source_width
  unless contiguous?(tensor)
    verify_storage_payload!(tensor)
    return write_strided_tensor(
      tensor, output, target_dtype: target_dtype, converter: converter,
                      elements_per_chunk: elements_per_chunk, source_width: source_width
    )
  end

  inline_payload = inline_storage_verification(tensor)
  verify_storage_payload!(tensor) unless inline_payload
  remaining = tensor.nbytes
  written = 0
  crc32 = 0 if inline_payload
  File.open(path, "rb") do |source|
    source.seek(tensor.data_start)
    while remaining.positive?
      requested = [remaining, bytes_per_chunk].min
      chunk = source.read(requested)
      raise Error, "Unexpected end of #{path} while reading #{tensor.name.inspect}" if chunk.nil? || chunk.bytesize != requested

      crc32 = Zlib.crc32(chunk, crc32) if inline_payload
      converted = converter.convert(chunk, from: tensor.dtype, to: target_dtype)
      output.write(converted)
      written += converted.bytesize
      remaining -= requested
    end
  end
  mark_inline_storage_verified!(tensor.storage_key, inline_payload, crc32) if inline_payload
  expected = tensor.element_count * Safetensors::DTYPE_BYTES.fetch(target_dtype)
  raise Error, "Converted tensor #{tensor.name.inspect} wrote #{written} bytes; expected #{expected}" unless written == expected

  written
end