jkrukowski/swift-safetensors

Installation

Add the following to your Package.swift file:

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

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

// get Swift Array (any numeric type)
let floatArray: [Float] = try parsedSafetensors.array(forKey: "tensorKey")
let int32Array: [Int32] = try parsedSafetensors.array(forKey: "tensorKey")
let doubleArray: [Double] = try parsedSafetensors.array(forKey: "tensorKey")

// get InlineArray (Swift 6.2+, requires macOS 26.0+)
let inlineArray: InlineArray<4, Float> = try parsedSafetensors.inlineArray(forKey: "tensorKey")

When MLTensor, MLMultiArray or MLShapedArray is materialized, the data is copied from the underlying buffer by default. If you want to avoid copying, you can use the noCopy parameter:

Note: Swift Array and InlineArray always own their data, so zero-copy access is not available for these types.

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

You can also use the SafetensorsBuilder to write Safetensors file:

import CoreML
import Safetensors

let builder = SafetensorsBuilder()
    .addTensor(MLShapedArray<Int32>(repeating: 1, shape: [2, 2]), forKey: "test1")
    .addTensor(MLShapedArray<Float32>(repeating: 2, shape: [9]), forKey: "test2")
    .withMetadata(["key1": "value1", "key2": "value2"])
    .withMaxShardingSize(1_000_000)

try builder.write(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

Package Metadata

Repository: jkrukowski/swift-safetensors

Default branch: main

README: README.md