defmodule Safetensors do
@moduledoc """
[Safetensors](https://huggingface.co/docs/safetensors/index) implementation for `Nx`.
## Examples
iex> x = Nx.tensor([1, 2, 3])
iex> y = Nx.tensor([1.0, 2.0, 3.0])
iex> tensors = %{"x" => x, "y" => y}
iex> data = Safetensors.dump(tensors)
iex> tensors = Safetensors.load!(data)
iex> tensors["x"]
#Nx.Tensor<
s64[3]
[1, 2, 3]
>
iex> tensors["y"]
#Nx.Tensor<
f32[3]
[1.0, 2.0, 3.0]
>
"""
alias Safetensors.Shared
@header_metadata_key "__metadata__"
@type_to_dtype %{
{:bf, 16} => "BF16",
{:f, 64} => "F64",
{:f, 32} => "F32",
{:f, 16} => "F16",
{:s, 64} => "I64",
{:s, 32} => "I32",
{:s, 16} => "I16",
{:s, 8} => "I8",
{:u, 64} => "U64",
{:u, 32} => "U32",
{:u, 16} => "U16",
{:u, 8} => "U8"
}
@dtype_to_type for {k, v} <- @type_to_dtype, into: %{}, do: {v, k}
@doc """
Writes a map of tensors to a file.
Tensors are written into the file one by one, without the need to
dump all of them into the memory at once.
"""
@spec write!(path :: Path.t(), %{String.t() => Nx.Tensor.t()}) :: :ok
def write!(path, tensors) when is_map(tensors) do
File.open!(path, [:write, :raw], fn file ->
tensors = Enum.sort(tensors)
{header_entries, _offset} =
Enum.map_reduce(tensors, 0, fn {tensor_name, tensor}, offset ->
tensor_header_entry(tensor_name, tensor, offset)
end)
:ok = :file.write(file, header_binary(header_entries))
for {_tensor_name, tensor} <- tensors do
:ok = :file.write(file, tensor_to_iodata(tensor))
end
end)
:ok
end
defp tensor_header_entry(tensor_name, tensor, offset) do
end_offset = offset + tensor_byte_size(tensor)
header_entry = {
tensor_name,
Jason.OrderedObject.new(
dtype: tensor |> Nx.type() |> type_to_dtype(),
shape: tensor |> Nx.shape() |> Tuple.to_list(),
data_offsets: [offset, end_offset]
)
}
{header_entry, end_offset}
end
defp header_binary(header_entries) do
header_json =
header_entries
|> Jason.OrderedObject.new()
|> Jason.encode!()
[<<byte_size(header_json)::unsigned-64-integer-little>>, header_json]
end
defp tensor_byte_size(tensor) do
{_, elem_size} = Nx.type(tensor)
elem_byte_size = div(elem_size, 8)
Nx.size(tensor) * elem_byte_size
end
defp tensor_to_iodata(tensor) do
{_, elem_size} = Nx.type(tensor)
tensor
|> Nx.to_binary()
|> Shared.new_byte_order(elem_size, :little)
end
@doc """
Serializes the given map of tensors to iodata.
`iodata` is a list of binaries that can be written to any IO device,
such as a file or a socket. You can ensure the result is a binary by
calling `IO.iodata_to_binary/1`.
"""
@spec dump(%{String.t() => Nx.Tensor.t()}) :: iodata()
def dump(tensors) when is_map(tensors) do
tensors = Enum.sort(tensors)
{header_entries, {buffer, _offset}} =
Enum.map_reduce(tensors, {[], 0}, fn {tensor_name, tensor}, {buffer, offset} ->
{header_entry, end_offset} = tensor_header_entry(tensor_name, tensor, offset)
binary = tensor_to_iodata(tensor)
{header_entry, {[buffer, binary], end_offset}}
end)
[header_binary(header_entries), buffer]
end
@doc """
Reads a serialized map of tensors from a file.
Tensors are loaded into Nx one by one, without the need to load the
entire file from disk into memory.
## Options
* `:lazy` - when `true`, instead of returning tensors, the function
returns lazy containers. Such a container can be converted to a
tensor using `Nx.to_tensor/1` and it is only at that point that
it is read from the file. Defaults to `false`
"""
@spec read!(path :: Path.t(), keyword()) :: %{String.t() => Nx.LazyContainer.t()}
def read!(path, opts \\ []) do
opts = Keyword.validate!(opts, lazy: false)
File.open!(path, [:read, :raw], fn file ->
{:ok, <<header_size::unsigned-64-integer-little>>} = :file.read(file, 8)
{:ok, header_json} = :file.read(file, header_size)
header = decode_header!(header_json)
for {tensor_name, tensor_info} <- header, into: %{} do
%{"data_offsets" => [offset_start, offset_end]} = tensor_info
{shape, type} = shape_and_type(tensor_info)
byte_offset = header_size + 8 + offset_start
byte_size = offset_end - offset_start
value =
if opts[:lazy] do
%Safetensors.FileTensor{
shape: shape,
type: type,
path: path,
byte_offset: byte_offset,
byte_size: byte_size
}
else
{:ok, binary} = :file.pread(file, byte_offset, byte_size)
Shared.build_tensor(binary, shape, type)
end
{tensor_name, value}
end
end)
end
@doc """
Loads a serialized map of tensors.
It is the opposite of `dump/1`.
"""
@spec load!(iodata()) :: %{String.t() => Nx.Tensor.t()}
def load!(data) when is_binary(data) or is_list(data) do
data = IO.iodata_to_binary(data)
<<
header_size::unsigned-64-integer-little,
header_json::binary-size(header_size),
buffer::binary
>> = data
header = decode_header!(header_json)
for {tensor_name, tensor_info} <- header, into: %{} do
%{"data_offsets" => [offset_start, offset_end]} = tensor_info
{shape, type} = shape_and_type(tensor_info)
tensor =
buffer
|> binary_slice(offset_start, offset_end - offset_start)
|> Shared.build_tensor(shape, type)
{tensor_name, tensor}
end
end
defp decode_header!(header_json) do
{_metadata, header} =
header_json
|> Jason.decode!()
|> Map.pop(@header_metadata_key)
header
end
defp shape_and_type(%{"dtype" => dtype, "shape" => shape}) do
{List.to_tuple(shape), dtype_to_type(dtype)}
end
defp type_to_dtype(type) do
@type_to_dtype[type] || raise "unrecognized type #{inspect(type)}"
end
defp dtype_to_type(dtype) do
@dtype_to_type[dtype] || raise "unrecognized dtype #{inspect(dtype)}"
end
end
