swift-safetensors

用于读取和写入 Safetensors 文件的 Swift 包。

安装

将以下内容添加到您的 Package.swift 文件中

dependencies: [
    .package(url: "https://github.com/jkrukowski/swift-safetensors", from: "0.0.6")
]

使用

读取 Safetensors 文件

import Safetensors

let parsedSafetensors = try Safetensors.read(
    at: URL(filePath: "path/to/file.safetensors")
)

// get MLTensor
let mlTensor = try parsedSafetensors.mlTensor(
    forKey: "tensorKey"
)

// get MLMultiArray
let mlMultiArray = try parsedSafetensors.mlMultiArray(
    forKey: "tensorKey"
)

// get MLShapedArray
let mlShapedArray: MLShapedArray<Int32> = try parsedSafetensors.mlShapedArray(
    forKey: "tensorKey"
)

当 MLTensor、MLMultiArray 或 MLShapedArray 被实例化时，数据会从底层缓冲区复制。 如果您想避免复制，您可以这样做

// get MLTensor without copying data
let mlTensor = try parsedSafetensors.mlTensor(
    forKey: "tensorKey",
    noCopy: true
)

// get MLMultiArray without copying data
let mlMultiArray = try parsedSafetensors.mlMultiArray(
    forKey: "tensorKey",
    noCopy: true
)

// get MLShapedArray without copying data
let mlShapedArray: MLShapedArray<Int32> = try parsedSafetensors.mlShapedArray(
    forKey: "tensorKey",
    noCopy: true
)

但请确保在您完成使用 MLTensor、MLMultiArray 或 MLShapedArray 之前，ParsedSafetensors 对象不会被释放。

写入 Safetensors 文件

import Safetensors

let data: [String: any SafetensorsEncodable] = [
    "test1": MLShapedArray<Int32>(repeating: 1, shape: [2, 2]),
    "test2": MLShapedArray<Float32>(repeating: 2, shape: [9]),
]

try Safetensors.write(
    data,
    metadata: ["key1": "value1", "key2": "value2"],
    to: URL(filePath: "path/to/file.safetensors")
)

代码格式化

本项目使用 swift-format。 要格式化代码，请运行

swift-format format . -i -r --configuration .swift-format