repugnant-pickle

因为它就是如此，难道不是吗？

这是一个用于处理 Python pickle 格式的 Rust 包。

它还支持打开 PyTorch 文件（.pth、.pt），它们基本上是包含 pickle 的 ZIP 文件。

什么？

Python pickle 格式。有些人赚得比我多得多，他们投入了大量的努力来开发版权保护方案，但这些方案在防止互操作性方面不如 Python pickle 格式有效。

为什么有人会真的创建这样的东西？我的理论是，创始人对他们的同类有着深深和持久的仇恨。他们唯一能得到任何乐趣的事情就是给尽可能多的人带来痛苦。为此：Python pickle 格式。

除非你是 Python，否则基本上无法可靠地处理。在一个理想的世界里，你只需让 pickle 独自存在即可，但遗憾的是，我们生活在这个世界上。有时你需要从一个 pickle 文件中获取一些数据，但你不想在应用程序中嵌入整个 Python 解释器或类似的东西。

这就是 repugnant-pickle 的用武之地。它尽力而为，并给你它所能提供的。

这个名字是对 BeautifulSoup（一个 HTML 抓取库）的戏仿。它没有（也不应该）试图完全解析 HTML（这是极其困难的），它只是给你一种相对简单的方式来提取所需的数据。大多数情况下，这已经足够了。

Example

以下是解析 PyTorch 文件时可能得到的输出示例

[Build(
  Global(
    Raw(GLOBAL("collections", "OrderedDict"), [
      Seq(Tuple, []),
      Seq(Tuple, [
        Seq(Dict, [
          String("emb.weight"),
          Global(Raw(GLOBAL("torch._utils", "_rebuild_tensor_v2")), [
            Seq(Tuple, [
              PersId(Seq(Tuple, [
                String("storage"),
                Raw(GLOBAL("torch", "BFloat16Storage")),
                String("0"),
                String("cuda:0"),
                Int(430348288),
              ])),
              Int(327378944),
              Seq(Tuple, [Int(50277), Int(1024)]),
              Seq(Tuple, [Int(1024), Int(1)]),
              Bool(false),
              Global(Raw(GLOBAL("collections", "OrderedDict")), [
                Seq(Tuple, [])
              ]),
            ]),
          ]),
        ]),
        Seq(Dict, [
          String("blocks.0.ln1.weight"),
          Global(Raw(GLOBAL("torch._utils", "_rebuild_tensor_v2")), [
            // Etc
          ]),
        ]),
      ]),
    ]),
  ),
)]

关于 serde-pickle 的看法？

如果你可以使用它，它会比这个包提供更好的体验。然而，它无法处理持久 ID 等一些事情。PyTorch 文件使用持久 ID，所以在这种情况下你真的没有选择。

还有一些奇怪的选项，比如强制 Rust 工具链使用 1.41 版本。

PyTorch

你可以启用 torch 功能以支持处理 PyTorch 文件。如果你的 Torch 文件包含奇怪的内容，你可能需要手动处理。查看 src/torch.rs 了解如何开始。

否则，您可以使用 repugnant_pickle::torch::RepugnantTorchTensors::new_from_file 来加载张量元数据。例如，如果它工作正常，您将得到类似以下内容

RepugnantTorchTensors(
   [
       RepugnantTorchTensor {
           name: "emb.weight",
           device: "cuda:0",
           tensor_type: BFloat16,
           storage: "archive/data/0",
           storage_len: 430348288,
           storage_offset: 327378944,
           absolute_offset: 327445248,
           shape: [50277, 1024],
           stride: [1024, 1],
           requires_grad: false,
       },
       RepugnantTorchTensor {
           name: "blocks.0.ln1.weight",
           device: "cuda:0",
           tensor_type: BFloat16,
           storage: "archive/data/0",
           storage_len: 430348288,
           storage_offset: 13639680,
           absolute_offset: 13705984,
           shape: [1024],
           stride: [1],
           requires_grad: false,
       },
   ]

您需要自己从类型和形状计算长度。例如，第一个张量的形状为 [1024, 50277]，类型为 BFloat16。因此，字节数为 (1024 * 50277) * 2，数据在 Torch 文件中的偏移量为 327445248，或者在 Torch ZIP 的 archive/data/0 文件中，它将从 327378944 开始。

cargo run --features torch --example dump_raw -- /path/to/torchfile.pth