Swift package for reading and writing Safetensors files.
Add the following to your Package.swift file:
dependencies: [
.package(url: "https://github.com/jkrukowski/swift-safetensors", from: "0.0.6")
]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"
)When MLTensor, MLMultiArray or MLShapedArray is materialized, the data is copied from the underlying buffer.
If you want to avoid copying, you can do:
// 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
)But make sure that the ParsedSafetensors object is not deallocated before you
finish using the MLTensor, MLMultiArray or MLShapedArray.
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")
)This project uses swift-format. To format the code run:
swift-format format . -i -r --configuration .swift-format