Class: Cohere::Transcribe::PyTorchCheckpoint::RestrictedUnpickler

Inherits:

Object

• Object
• Cohere::Transcribe::PyTorchCheckpoint::RestrictedUnpickler

Defined in:

lib/cohere/transcribe/pytorch_checkpoint.rb

Overview

Restricted pickle virtual machine for tensor state dictionaries.

Constant Summary

MARKER =

Object.new.freeze

Instance Attribute Summary

• #storages ⇒ Object readonly

  Returns the value of attribute storages.

Instance Method Summary

• #initialize(io, limit: PICKLE_LIMIT, storages: nil) ⇒ RestrictedUnpickler constructor

  A new instance of RestrictedUnpickler.

• #load ⇒ Object

Constructor Details

#initialize(io, limit: PICKLE_LIMIT, storages: nil) ⇒ RestrictedUnpickler

Returns a new instance of RestrictedUnpickler.

# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 386

def initialize(io, limit: PICKLE_LIMIT, storages: nil)
  @io = io
  @limit = limit
  @start = io.pos
  @stack = []
  @memo = {}
  @storages = storages || {}
  @next_memo = 0
  @opcode_count = 0
end

Instance Attribute Details

#storages ⇒ Object (readonly)

Returns the value of attribute storages.

# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 384

def storages
  @storages
end

Instance Method Details

#load ⇒ Object

# File 'lib/cohere/transcribe/pytorch_checkpoint.rb', line 397

def load
  loop do
    @opcode_count += 1
    raise Error, "Restricted PyTorch pickle exceeds the opcode-count limit" if @opcode_count > PICKLE_OPCODE_LIMIT

    opcode = read_exact(1).getbyte(0)
    case opcode
    when 0x80 then protocol!
    when 0x95 then read_uint64 # FRAME
    when 0x2e then return pop # STOP
    when 0x28 then push(MARKER)
    when 0x30 then pop
    when 0x31 then pop_mark
    when 0x32 then push(peek)
    when 0x4e then push(nil)
    when 0x88 then push(true)
    when 0x89 then push(false)
    when 0x4a then push(read_exact(4).unpack1("l<"))
    when 0x4b then push(read_exact(1).unpack1("C"))
    when 0x4d then push(read_exact(2).unpack1("v"))
    when 0x49 then push(parse_ascii_integer(read_line))
    when 0x4c then push(Integer(read_line.delete_suffix("L"), 10))
    when 0x8a then push(parse_long(read_long_bytes(1)))
    when 0x8b then push(parse_long(read_long_bytes(4)))
    when 0x46 then push(Float(read_line))
    when 0x47 then push(read_exact(8).unpack1("G"))
    when 0x58 then push(read_utf8(read_exact(read_exact(4).unpack1("V"))))
    when 0x8c then push(read_utf8(read_exact(read_exact(1).unpack1("C"))))
    when 0x8d then push(read_utf8(read_exact(read_uint64)))
    when 0x54 then push(read_exact(read_exact(4).unpack1("V")))
    when 0x55 then push(read_exact(read_exact(1).unpack1("C")))
    when 0x42 then push(read_exact(read_exact(4).unpack1("V")))
    when 0x43 then push(read_exact(read_exact(1).unpack1("C")))
    when 0x8e then push(read_exact(read_uint64))
    when 0x96 then push(read_exact(read_uint64).dup)
    when 0x53 then push(parse_quoted_string(read_line))
    when 0x5d then push([])
    when 0x6c then push(pop_mark)
    when 0x61
      value = pop
      list_target << value
    when 0x65
      values = pop_mark
      list_target.concat(values)
    when 0x7d then push({})
    when 0x64 then push(Hash[*pop_mark])
    when 0x73 then set_pair
    when 0x75 then set_pairs
    when 0x29 then push([])
    when 0x74 then push(pop_mark)
    when 0x85 then push([pop])
    when 0x86 then push(pop(2))
    when 0x87 then push(pop(3))
    when 0x8f then push([]) # EMPTY_SET; arrays suffice for metadata
    when 0x90
      values = pop_mark
      list_target.concat(values).uniq!
    when 0x91 then push(pop_mark.uniq.freeze)
    when 0x63 then push(read_global)
    when 0x93
      module_name, name = pop(2)
      raise Error, "Restricted PyTorch pickle has an invalid STACK_GLOBAL" unless module_name.is_a?(String) && name.is_a?(String)

      push(Global.new(module_name: module_name, name: name))
    when 0x52 then reduce!
    when 0x51 then push(persistent_load(pop))
    when 0x50 then push(persistent_load(read_line))
    when 0x62 then build!
    when 0x81 then new_object!
    when 0x92 then new_object_with_keywords!
    when 0x71 then memo_store(read_exact(1).unpack1("C"))
    when 0x72 then memo_store(read_exact(4).unpack1("V"))
    when 0x70 then memo_store(Integer(read_line, 10))
    when 0x94 then memo_store_next
    when 0x68 then push(memo_fetch(read_exact(1).unpack1("C")))
    when 0x6a then push(memo_fetch(read_exact(4).unpack1("V")))
    when 0x67 then push(memo_fetch(Integer(read_line, 10)))
    else
      raise Error, format("Restricted PyTorch pickle uses unsupported opcode 0x%02x", opcode)
    end
  end
rescue EOFError
  raise Error, "PyTorch pickle ended before STOP"
rescue ArgumentError, TypeError, RangeError => e
  raise Error, "Invalid restricted PyTorch pickle: #{e.message}"
end