GGUF: Load and save metadata (#446)

* gguf metadata
---------

Co-authored-by: Awni Hannun <awni@apple.com>
This commit is contained in:
Juarez Bochi
2024-01-19 23:06:05 +01:00
committed by GitHub
parent 6589c869d6
commit ddf50113c5
11 changed files with 668 additions and 94 deletions

View File

@@ -7,26 +7,36 @@
#include <string>
#include <unordered_map>
#include <variant>
#include "mlx/ops.h"
#include "mlx/io.h"
namespace py = pybind11;
using namespace mlx::core;
using DictOrArray = std::variant<array, std::unordered_map<std::string, array>>;
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>>>;
std::unordered_map<std::string, array> mlx_load_safetensor_helper(
py::object file,
StreamOrDevice s);
void mlx_save_safetensor_helper(py::object file, py::dict d);
std::unordered_map<std::string, array> mlx_load_gguf_helper(
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,
StreamOrDevice s);
void mlx_save_gguf_helper(py::object file, py::dict d);
py::dict d,
std::optional<py::dict> m);
DictOrArray mlx_load_helper(
LoadOutputTypes mlx_load_helper(
py::object file,
std::optional<std::string> format,
bool return_metadata,
StreamOrDevice s);
void mlx_save_helper(py::object file, array a);
void mlx_savez_helper(