Metadata support for safetensors (#639)

* metadata support for safetensors

* aliases making it alittle more readable

* addressing comments

* python binding tests
This commit is contained in:
Diogo
2024-02-08 22:33:15 -05:00
committed by GitHub
parent 221f8d3fc2
commit b57bd0488d
8 changed files with 108 additions and 69 deletions

View File

@@ -15,19 +15,17 @@ using namespace mlx::core;
using LoadOutputTypes = std::variant<
array,
std::unordered_map<std::string, array>,
std::pair<
std::unordered_map<std::string, array>,
std::unordered_map<std::string, MetaData>>>;
SafetensorsLoad,
GGUFLoad>;
std::unordered_map<std::string, array> mlx_load_safetensor_helper(
SafetensorsLoad mlx_load_safetensor_helper(py::object file, StreamOrDevice s);
void mlx_save_safetensor_helper(
py::object file,
StreamOrDevice s);
void mlx_save_safetensor_helper(py::object file, py::dict d);
py::dict d,
std::optional<py::dict> m);
GGUFLoad mlx_load_gguf_helper(py::object file, StreamOrDevice s);
std::pair<
std::unordered_map<std::string, array>,
std::unordered_map<std::string, MetaData>>
mlx_load_gguf_helper(py::object file, StreamOrDevice s);
void mlx_save_gguf_helper(
py::object file,
py::dict d,