swift-safetensors

Swift package for reading and writing Safetensors files.

Installation

Add the following to your Package.swift file:

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

Usage

Read Safetensors file

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.

Write Safetensors file

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")
)

Code Formatting

This project uses swift-format. To format the code run:

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