Class: Cohere::Transcribe::PyTorchCheckpoint::ZipArchive

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

Overview

Minimal ZIP/ZIP64 central-directory reader. Torch stores tensor records without compression, allowing bounded direct reads from multi-GB files.

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(path) ⇒ ZipArchive

Returns a new instance of ZipArchive.



67
68
69
70
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 67

def initialize(path)
  @path = Pathname(path).expand_path
  @entries = read_entries.freeze
end

Instance Attribute Details

#entriesObject (readonly)

Returns the value of attribute entries.



65
66
67
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 65

def entries
  @entries
end

#pathObject (readonly)

Returns the value of attribute path.



65
66
67
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 65

def path
  @path
end

Instance Method Details

#fetch(name) ⇒ Object



72
73
74
75
76
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 72

def fetch(name)
  entries.fetch(name)
rescue KeyError
  raise Error, "PyTorch archive #{path} has no entry #{name.inspect}"
end

#read(name, limit: PICKLE_LIMIT) ⇒ Object



78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 78

def read(name, limit: PICKLE_LIMIT)
  entry = fetch(name)
  raise Error, "PyTorch archive entry #{name.inspect} exceeds the #{limit}-byte size limit" if entry.size > limit
  if entry.compressed_size > limit + ZIP_COMPRESSED_OVERHEAD
    raise Error,
          "PyTorch archive entry #{name.inspect} has an oversized compressed representation"
  end
  if entry.compression_method.zero? && entry.compressed_size != entry.size
    raise Error, "Stored PyTorch archive entry #{name.inspect} has inconsistent sizes"
  end

  compressed = read_range(entry.data_offset, entry.compressed_size)
  payload = case entry.compression_method
            when 0
              compressed
            when 8
              inflate_bounded(compressed, entry, limit)
            else
              raise Error,
                    "PyTorch archive entry #{name.inspect} uses unsupported ZIP method " \
                    "#{entry.compression_method}"
            end
  if payload.bytesize != entry.size
    raise Error,
          "PyTorch archive entry #{name.inspect} decoded to #{payload.bytesize} bytes; expected #{entry.size}"
  end
  raise Error, "PyTorch archive entry #{name.inspect} failed its CRC check" unless Zlib.crc32(payload) == entry.crc32

  payload
rescue Zlib::Error => e
  raise Error, "Cannot inflate PyTorch archive entry #{name.inspect}: #{e.message}"
end