-module(mlx_io). -include_lib("kernel/include/file.hrl"). %% Model serialization, checkpointing, and format conversion -export([ %% Model saving and loading save_model/2, save_model/3, load_model/1, load_model/2, %% Checkpointing save_checkpoint/2, save_checkpoint/3, load_checkpoint/1, load_checkpoint/2, list_checkpoints/1, cleanup_checkpoints/2, %% Format conversion convert_format/3, convert_format/4, supported_formats/0, %% Safetensors support save_safetensors/2, load_safetensors/1, verify_safetensors/1, %% ONNX support export_onnx/2, export_onnx/3, import_onnx/1, import_onnx/2, %% PyTorch/HuggingFace compatibility save_pytorch_state_dict/2, load_pytorch_state_dict/1, convert_from_huggingface/2, %% TensorFlow compatibility export_tensorflow/2, export_tensorflow/3, import_tensorflow/1, import_tensorflow/2, %% Quantized model I/O save_quantized_model/3, load_quantized_model/1, %% Model metadata save_metadata/2, load_metadata/1, extract_model_info/1, %% Distributed checkpointing save_distributed_checkpoint/3, load_distributed_checkpoint/2, merge_distributed_checkpoints/2, %% Model compression for storage compress_for_storage/2, compress_for_storage/3, decompress_from_storage/1, %% Streaming I/O for large models create_model_stream/2, stream_model_weights/2, close_model_stream/1 ]). %% Model saving and loading save_model(Model, FilePath) -> save_model(Model, FilePath, #{}). save_model(Model, FilePath, Options) -> % Save model in MLX native format Format = maps:get(format, Options, mlx_native), Compression = maps:get(compression, Options, none), IncludeMetadata = maps:get(include_metadata, Options, true), % Extract model components Weights = extract_weights(Model), Architecture = extract_architecture(Model), Metadata = case IncludeMetadata of true -> extract_model_metadata(Model); false -> #{} end, % Prepare data for serialization ModelData = #{ format_version => "1.0", architecture => Architecture, weights => Weights, metadata => Metadata, creation_time => erlang:system_time(second) }, % Serialize and compress if needed SerializedData = serialize_model_data(ModelData, Format), FinalData = case Compression of none -> SerializedData; gzip -> zlib:gzip(SerializedData); lz4 -> compress_lz4(SerializedData) end, % Write to file case file:write_file(FilePath, FinalData) of ok -> FileSize = byte_size(FinalData), {ok, #{file_size => FileSize, format => Format, compression => Compression}}; {error, Reason} -> {error, {file_write_failed, Reason}} end. load_model(FilePath) -> load_model(FilePath, #{}). load_model(FilePath, Options) -> % Load model from file case file:read_file(FilePath) of {ok, FileData} -> % Detect compression and decompress if needed DecompressedData = detect_and_decompress(FileData), % Deserialize model data case deserialize_model_data(DecompressedData) of {ok, ModelData} -> % Reconstruct model Architecture = maps:get(architecture, ModelData), Weights = maps:get(weights, ModelData), reconstruct_model(Architecture, Weights, Options); {error, Reason} -> {error, {deserialization_failed, Reason}} end; {error, Reason} -> {error, {file_read_failed, Reason}} end. %% Checkpointing save_checkpoint(Model, CheckpointDir) -> save_checkpoint(Model, CheckpointDir, #{}). save_checkpoint(Model, CheckpointDir, Options) -> % Save training checkpoint Epoch = maps:get(epoch, Options, 0), Step = maps:get(step, Options, 0), OptimizerState = maps:get(optimizer_state, Options, #{}), TrainingState = maps:get(training_state, Options, #{}), % Create checkpoint filename Timestamp = erlang:system_time(second), CheckpointName = io_lib:format("checkpoint_epoch_~w_step_~w_~w.ckpt", [Epoch, Step, Timestamp]), CheckpointPath = filename:join(CheckpointDir, CheckpointName), % Prepare checkpoint data CheckpointData = #{ model => extract_weights(Model), optimizer_state => OptimizerState, training_state => TrainingState, epoch => Epoch, step => Step, timestamp => Timestamp, mlx_version => get_mlx_version() }, % Ensure checkpoint directory exists ok = filelib:ensure_dir(CheckpointPath), % Save checkpoint case save_model_data(CheckpointData, CheckpointPath) of ok -> % Update latest checkpoint symlink LatestPath = filename:join(CheckpointDir, "latest.ckpt"), case file:make_symlink(CheckpointName, LatestPath) of ok -> {ok, CheckpointPath}; {error, eexist} -> file:delete(LatestPath), file:make_symlink(CheckpointName, LatestPath), {ok, CheckpointPath}; {error, Reason} -> {error, {symlink_failed, Reason}} end; {error, Reason} -> {error, Reason} end. load_checkpoint(CheckpointPath) -> load_checkpoint(CheckpointPath, #{}). load_checkpoint(CheckpointPath, Options) -> % Load training checkpoint LoadLatest = maps:get(load_latest, Options, false), ActualPath = case LoadLatest of true -> CheckpointDir = filename:dirname(CheckpointPath), LatestPath = filename:join(CheckpointDir, "latest.ckpt"), case file:read_link(LatestPath) of {ok, LatestFile} -> filename:join(CheckpointDir, LatestFile); {error, _} -> CheckpointPath end; false -> CheckpointPath end, case load_model_data(ActualPath) of {ok, CheckpointData} -> Model = maps:get(model, CheckpointData), OptimizerState = maps:get(optimizer_state, CheckpointData, #{}), TrainingState = maps:get(training_state, CheckpointData, #{}), {ok, #{ model => Model, optimizer_state => OptimizerState, training_state => TrainingState, epoch => maps:get(epoch, CheckpointData, 0), step => maps:get(step, CheckpointData, 0) }}; {error, Reason} -> {error, Reason} end. list_checkpoints(CheckpointDir) -> % List all checkpoints in directory Pattern = filename:join(CheckpointDir, "checkpoint_*.ckpt"), case filelib:wildcard(Pattern) of [] -> {ok, []}; Files -> % Sort by modification time FilesWithTime = lists:map(fun(File) -> case file:read_file_info(File) of {ok, FileInfo} -> {File, FileInfo#file_info.mtime}; {error, _} -> {File, 0} end end, Files), SortedFiles = lists:reverse(lists:keysort(2, FilesWithTime)), {ok, [File || {File, _} <- SortedFiles]} end. cleanup_checkpoints(CheckpointDir, KeepLast) -> % Keep only the last N checkpoints case list_checkpoints(CheckpointDir) of {ok, Checkpoints} when length(Checkpoints) > KeepLast -> ToDelete = lists:nthtail(KeepLast, Checkpoints), DeleteResults = lists:map(fun(File) -> case file:delete(File) of ok -> {ok, File}; {error, Reason} -> {error, {File, Reason}} end end, ToDelete), {ok, DeleteResults}; {ok, _} -> {ok, []}; {error, Reason} -> {error, Reason} end. %% Format conversion convert_format(InputPath, OutputPath, TargetFormat) -> convert_format(InputPath, OutputPath, TargetFormat, #{}). convert_format(InputPath, OutputPath, TargetFormat, Options) -> % Convert between different model formats case load_model(InputPath) of {ok, Model} -> case TargetFormat of onnx -> export_onnx(Model, OutputPath, Options); tensorflow -> export_tensorflow(Model, OutputPath, Options); pytorch -> save_pytorch_state_dict(Model, OutputPath); safetensors -> save_safetensors(Model, OutputPath); mlx_native -> save_model(Model, OutputPath, Options); _ -> {error, {unsupported_format, TargetFormat}} end; {error, Reason} -> {error, Reason} end. supported_formats() -> [mlx_native, onnx, tensorflow, pytorch, safetensors]. %% Safetensors support save_safetensors(Model, FilePath) -> % Save model in Safetensors format Weights = extract_weights(Model), % Convert weights to safetensors format SafetensorsData = convert_to_safetensors(Weights), % Write safetensors file mlx_nif:save_safetensors(SafetensorsData, FilePath). load_safetensors(FilePath) -> % Load model from Safetensors format case mlx_nif:load_safetensors(FilePath) of {ok, WeightsData} -> Weights = convert_from_safetensors(WeightsData), {ok, Weights}; {error, Reason} -> {error, Reason} end. verify_safetensors(FilePath) -> % Verify safetensors file integrity mlx_nif:verify_safetensors(FilePath). %% ONNX support export_onnx(Model, FilePath) -> export_onnx(Model, FilePath, #{}). export_onnx(Model, FilePath, Options) -> % Export model to ONNX format OpsetVersion = maps:get(opset_version, Options, 17), InputShape = maps:get(input_shape, Options, undefined), InputNames = maps:get(input_names, Options, ["input"]), OutputNames = maps:get(output_names, Options, ["output"]), % Convert MLX model to ONNX graph case convert_to_onnx_graph(Model, InputShape, OpsetVersion) of {ok, ONNXGraph} -> % Create ONNX model proto ONNXModel = create_onnx_model(ONNXGraph, InputNames, OutputNames, OpsetVersion), % Serialize and save case serialize_onnx_model(ONNXModel) of {ok, SerializedData} -> file:write_file(FilePath, SerializedData); {error, Reason} -> {error, {onnx_serialization_failed, Reason}} end; {error, Reason} -> {error, {onnx_conversion_failed, Reason}} end. import_onnx(FilePath) -> import_onnx(FilePath, #{}). import_onnx(FilePath, Options) -> % Import ONNX model to MLX case file:read_file(FilePath) of {ok, ONNXData} -> case deserialize_onnx_model(ONNXData) of {ok, ONNXModel} -> convert_from_onnx_model(ONNXModel, Options); {error, Reason} -> {error, {onnx_deserialization_failed, Reason}} end; {error, Reason} -> {error, {file_read_failed, Reason}} end. %% PyTorch/HuggingFace compatibility save_pytorch_state_dict(Model, FilePath) -> % Save model in PyTorch state_dict format Weights = extract_weights(Model), % Convert to PyTorch tensor format PytorchStateDict = convert_to_pytorch_format(Weights), % Save as pickle file (simplified - would need proper Python pickle format) SerializedData = serialize_pytorch_state_dict(PytorchStateDict), file:write_file(FilePath, SerializedData). load_pytorch_state_dict(FilePath) -> % Load PyTorch state_dict case file:read_file(FilePath) of {ok, PickleData} -> case deserialize_pytorch_state_dict(PickleData) of {ok, StateDict} -> Weights = convert_from_pytorch_format(StateDict), {ok, Weights}; {error, Reason} -> {error, Reason} end; {error, Reason} -> {error, {file_read_failed, Reason}} end. convert_from_huggingface(ModelName, CacheDir) -> % Convert HuggingFace model to MLX format % This would typically involve downloading from HuggingFace Hub % and converting the model weights and configuration ConfigPath = filename:join([CacheDir, ModelName, "config.json"]), WeightsPath = filename:join([CacheDir, ModelName, "pytorch_model.bin"]), case {file:read_file(ConfigPath), load_pytorch_state_dict(WeightsPath)} of {{ok, ConfigData}, {ok, Weights}} -> Config = jsx:decode(ConfigData, [return_maps]), Architecture = convert_hf_config_to_mlx(Config), Model = reconstruct_model(Architecture, Weights, #{}), {ok, Model}; {{error, Reason}, _} -> {error, {config_load_failed, Reason}}; {_, {error, Reason}} -> {error, {weights_load_failed, Reason}} end. %% TensorFlow compatibility export_tensorflow(Model, FilePath) -> export_tensorflow(Model, FilePath, #{}). export_tensorflow(Model, FilePath, Options) -> % Export model to TensorFlow SavedModel format SignatureName = maps:get(signature_name, Options, "serving_default"), % Convert MLX model to TensorFlow graph case convert_to_tf_graph(Model) of {ok, TFGraph} -> % Create SavedModel SavedModel = create_tf_saved_model(TFGraph, SignatureName), % Save to directory case file:make_dir(FilePath) of ok -> save_tf_saved_model(SavedModel, FilePath); {error, eexist} -> save_tf_saved_model(SavedModel, FilePath); {error, Reason} -> {error, {dir_creation_failed, Reason}} end; {error, Reason} -> {error, {tf_conversion_failed, Reason}} end. import_tensorflow(SavedModelPath) -> import_tensorflow(SavedModelPath, #{}). import_tensorflow(SavedModelPath, Options) -> % Import TensorFlow SavedModel case load_tf_saved_model(SavedModelPath) of {ok, SavedModel} -> convert_from_tf_saved_model(SavedModel, Options); {error, Reason} -> {error, Reason} end. %% Quantized model I/O save_quantized_model(Model, QuantParams, FilePath) -> % Save quantized model with quantization parameters ModelWeights = extract_weights(Model), QuantizedData = #{ quantized_weights => ModelWeights, quantization_params => QuantParams, quantization_version => "1.0", creation_time => erlang:system_time(second) }, save_model_data(QuantizedData, FilePath). load_quantized_model(FilePath) -> % Load quantized model case load_model_data(FilePath) of {ok, QuantizedData} -> Weights = maps:get(quantized_weights, QuantizedData), QuantParams = maps:get(quantization_params, QuantizedData), {ok, {Weights, QuantParams}}; {error, Reason} -> {error, Reason} end. %% Model metadata save_metadata(Metadata, FilePath) -> % Save model metadata to JSON MetadataJson = jsx:encode(Metadata), file:write_file(FilePath, MetadataJson). load_metadata(FilePath) -> % Load model metadata from JSON case file:read_file(FilePath) of {ok, MetadataData} -> try Metadata = jsx:decode(MetadataData, [return_maps]), {ok, Metadata} catch error:Reason -> {error, {json_decode_failed, Reason}} end; {error, Reason} -> {error, {file_read_failed, Reason}} end. extract_model_info(Model) -> % Extract comprehensive model information Weights = extract_weights(Model), % Calculate model statistics TotalParams = maps:fold(fun(_Name, Weight, Acc) -> {ok, Size} = mlx:size(Weight), Acc + Size end, 0, Weights), ModelSize = maps:fold(fun(_Name, Weight, Acc) -> {ok, Size} = mlx:size(Weight), {ok, Dtype} = mlx:dtype(Weight), BytesPerElement = case Dtype of float32 -> 4; float16 -> 2; int8 -> 1; int4 -> 0.5 end, Acc + Size * BytesPerElement end, 0, Weights), LayerInfo = maps:map(fun(_Name, Weight) -> {ok, Shape} = mlx:shape(Weight), {ok, Dtype} = mlx:dtype(Weight), {ok, Size} = mlx:size(Weight), #{shape => Shape, dtype => Dtype, size => Size} end, Weights), #{ total_parameters => TotalParams, model_size_bytes => ModelSize, num_layers => maps:size(Weights), layer_info => LayerInfo, creation_time => erlang:system_time(second) }. %% Distributed checkpointing save_distributed_checkpoint(Model, CheckpointDir, NodeId) -> % Save checkpoint for distributed training NodeCheckpointDir = filename:join(CheckpointDir, io_lib:format("node_~w", [NodeId])), ok = filelib:ensure_dir(filename:join(NodeCheckpointDir, "dummy")), % Save local model shard LocalWeights = extract_local_weights(Model, NodeId), CheckpointPath = filename:join(NodeCheckpointDir, "model_shard.ckpt"), case save_model_data(LocalWeights, CheckpointPath) of ok -> % Save metadata about this shard ShardMetadata = #{ node_id => NodeId, shard_keys => maps:keys(LocalWeights), timestamp => erlang:system_time(second) }, MetadataPath = filename:join(NodeCheckpointDir, "shard_metadata.json"), save_metadata(ShardMetadata, MetadataPath); {error, Reason} -> {error, Reason} end. load_distributed_checkpoint(CheckpointDir, NodeId) -> % Load checkpoint for specific node NodeCheckpointDir = filename:join(CheckpointDir, io_lib:format("node_~w", [NodeId])), CheckpointPath = filename:join(NodeCheckpointDir, "model_shard.ckpt"), load_model_data(CheckpointPath). merge_distributed_checkpoints(CheckpointDir, OutputPath) -> % Merge distributed checkpoints into single model % Find all node directories Pattern = filename:join(CheckpointDir, "node_*"), NodeDirs = filelib:wildcard(Pattern), % Load all shards AllShards = lists:foldl(fun(NodeDir, Acc) -> CheckpointPath = filename:join(NodeDir, "model_shard.ckpt"), case load_model_data(CheckpointPath) of {ok, Shard} -> maps:merge(Acc, Shard); {error, _} -> Acc end end, #{}, NodeDirs), % Save merged model save_model_data(AllShards, OutputPath). %% Model compression for storage compress_for_storage(Model, CompressionLevel) -> compress_for_storage(Model, CompressionLevel, #{}). compress_for_storage(Model, CompressionLevel, Options) -> % Compress model for efficient storage Method = maps:get(method, Options, adaptive), Weights = extract_weights(Model), case Method of lossless -> % Use lossless compression (gzip, lz4, etc.) SerializedWeights = serialize_weights(Weights), CompressedData = compress_data(SerializedWeights, CompressionLevel), {ok, CompressedData}; lossy -> % Use lossy compression (quantization + compression) QuantizedWeights = apply_storage_quantization(Weights, CompressionLevel), SerializedWeights = serialize_weights(QuantizedWeights), CompressedData = compress_data(SerializedWeights, CompressionLevel), {ok, CompressedData}; adaptive -> % Choose compression method per layer CompressedLayers = maps:map(fun(LayerName, Weight) -> choose_optimal_compression(LayerName, Weight, CompressionLevel) end, Weights), {ok, CompressedLayers} end. decompress_from_storage(CompressedModel) -> % Decompress model from storage format case detect_compression_format(CompressedModel) of {ok, Format} -> decompress_with_format(CompressedModel, Format); {error, Reason} -> {error, Reason} end. %% Streaming I/O for large models create_model_stream(FilePath, Mode) -> % Create streaming interface for large models case Mode of read -> case file:open(FilePath, [read, binary]) of {ok, FileHandle} -> StreamState = #{ file_handle => FileHandle, mode => read, position => 0, buffer_size => 1024 * 1024 % 1MB buffer }, {ok, StreamState}; {error, Reason} -> {error, Reason} end; write -> case file:open(FilePath, [write, binary]) of {ok, FileHandle} -> StreamState = #{ file_handle => FileHandle, mode => write, position => 0, buffer => <<>> }, {ok, StreamState}; {error, Reason} -> {error, Reason} end end. stream_model_weights(StreamState, LayerName) -> % Stream individual layer weights case maps:get(mode, StreamState) of read -> FileHandle = maps:get(file_handle, StreamState), BufferSize = maps:get(buffer_size, StreamState), case read_layer_from_stream(FileHandle, LayerName, BufferSize) of {ok, WeightData, NewPosition} -> NewState = maps:put(position, NewPosition, StreamState), {ok, WeightData, NewState}; {error, Reason} -> {error, Reason} end; write -> % Implementation for writing weights to stream {error, not_implemented} end. close_model_stream(StreamState) -> % Close model stream FileHandle = maps:get(file_handle, StreamState), file:close(FileHandle). %% Helper functions extract_weights(Model) -> % Extract weights from model (simplified) Model. extract_architecture(Model) -> % Extract architecture description (simplified) #{type => unknown}. extract_model_metadata(Model) -> % Extract model metadata (simplified) #{}. serialize_model_data(ModelData, Format) -> % Serialize model data in specified format case Format of mlx_native -> term_to_binary(ModelData); json -> jsx:encode(ModelData); _ -> term_to_binary(ModelData) end. deserialize_model_data(SerializedData) -> % Deserialize model data try ModelData = binary_to_term(SerializedData), {ok, ModelData} catch error:Reason -> {error, Reason} end. reconstruct_model(Architecture, Weights, Options) -> % Reconstruct model from architecture and weights {ok, Weights}. % Simplified detect_and_decompress(FileData) -> % Detect compression and decompress case FileData of <<31, 139, _/binary>> -> % gzip magic number zlib:gunzip(FileData); _ -> FileData end. compress_lz4(Data) -> % LZ4 compression (would use proper LZ4 library) Data. save_model_data(Data, FilePath) -> SerializedData = term_to_binary(Data), file:write_file(FilePath, SerializedData). load_model_data(FilePath) -> case file:read_file(FilePath) of {ok, SerializedData} -> try Data = binary_to_term(SerializedData), {ok, Data} catch error:Reason -> {error, {deserialization_failed, Reason}} end; {error, Reason} -> {error, {file_read_failed, Reason}} end. get_mlx_version() -> "1.0.0". convert_to_safetensors(Weights) -> % Convert to safetensors format Weights. convert_from_safetensors(WeightsData) -> % Convert from safetensors format WeightsData. convert_to_onnx_graph(Model, InputShape, OpsetVersion) -> % Convert MLX model to ONNX graph {ok, #{}}. create_onnx_model(Graph, InputNames, OutputNames, OpsetVersion) -> % Create ONNX model proto #{}. serialize_onnx_model(ONNXModel) -> % Serialize ONNX model {ok, <<>>}. deserialize_onnx_model(ONNXData) -> % Deserialize ONNX model {ok, #{}}. convert_from_onnx_model(ONNXModel, Options) -> % Convert ONNX model to MLX {ok, #{}}. convert_to_pytorch_format(Weights) -> % Convert to PyTorch format Weights. serialize_pytorch_state_dict(StateDict) -> % Serialize PyTorch state dict term_to_binary(StateDict). deserialize_pytorch_state_dict(PickleData) -> % Deserialize PyTorch state dict try StateDict = binary_to_term(PickleData), {ok, StateDict} catch error:Reason -> {error, Reason} end. convert_from_pytorch_format(StateDict) -> % Convert from PyTorch format StateDict. convert_hf_config_to_mlx(Config) -> % Convert HuggingFace config to MLX architecture Config. convert_to_tf_graph(Model) -> % Convert MLX model to TensorFlow graph {ok, #{}}. create_tf_saved_model(TFGraph, SignatureName) -> % Create TensorFlow SavedModel #{}. save_tf_saved_model(SavedModel, FilePath) -> % Save TensorFlow SavedModel ok. load_tf_saved_model(SavedModelPath) -> % Load TensorFlow SavedModel {ok, #{}}. convert_from_tf_saved_model(SavedModel, Options) -> % Convert TensorFlow SavedModel to MLX {ok, #{}}. extract_local_weights(Model, NodeId) -> % Extract weights for specific node in distributed setup Model. serialize_weights(Weights) -> term_to_binary(Weights). compress_data(Data, CompressionLevel) -> zlib:gzip(Data). apply_storage_quantization(Weights, CompressionLevel) -> % Apply quantization for storage Weights. choose_optimal_compression(LayerName, Weight, CompressionLevel) -> % Choose optimal compression for layer Weight. detect_compression_format(CompressedModel) -> % Detect compression format {ok, gzip}. decompress_with_format(CompressedModel, Format) -> % Decompress with specific format case Format of gzip -> zlib:gunzip(CompressedModel); _ -> CompressedModel end. read_layer_from_stream(FileHandle, LayerName, BufferSize) -> % Read layer from file stream case file:read(FileHandle, BufferSize) of {ok, Data} -> {ok, Data, file:position(FileHandle, cur)}; eof -> {error, eof}; {error, Reason} -> {error, Reason} end.