Module: Cohere::Transcribe::PyTorchCheckpoint

Defined in:

lib/cohere/transcribe/pytorch_checkpoint.rb

Overview

A non-executing reader for dense PyTorch state dictionaries. It accepts the ZIP serialization used by current torch.save and the immediately preceding pickle-plus-storage stream. Pickle GLOBAL/REDUCE operations are interpreted through a small allowlist; no Python classes or code run.

Defined Under Namespace

Classes: Entry, Error, Global, Reader, RestrictedUnpickler, StorageRef, Tensor, TensorSet, TensorSpec, ZipArchive

Constant Summary

MAGIC_NUMBER =

119_547_037_146_038_801_333_356

PROTOCOL_VERSION =

1001

PICKLE_LIMIT =

64 * 1024 * 1024

PICKLE_OPCODE_LIMIT =

1_000_000

PICKLE_LONG_BYTES_LIMIT =

16

ZIP_TAIL_LIMIT =

65_557

ZIP_ENTRY_LIMIT =

1_000_000

ZIP_CENTRAL_LIMIT =

256 * 1024 * 1024

ZIP_COMPRESSED_OVERHEAD =

1024 * 1024

ZIP_INFLATE_CHUNK =

16 * 1024

ZIP_EOCD =

0x0605_4b50

ZIP64_EOCD =

0x0606_4b50

ZIP64_LOCATOR =

0x0706_4b50

ZIP_CENTRAL =

0x0201_4b50

ZIP_LOCAL =

0x0403_4b50

STORAGE_DTYPES =

{
  "FloatStorage" => ["F32", 4],
  "HalfStorage" => ["F16", 2],
  "BFloat16Storage" => ["BF16", 2],
  "LongStorage" => ["I64", 8]
}.freeze

FLOAT_DTYPES =

%w[BF16 F16 F32].freeze