zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-08-21 20:46:46 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

covering more cases

Browse Source
This commit is contained in:
dc-dc-dc 2023-12-17 21:13:10 -05:00
parent 87ec7b3cf9
commit bd422decc4
3 changed files with 174 additions and 71 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

168
mlx/safetensor.cpp
View File

@ -36,17 +36,14 @@ Token Tokenizer::getToken() {
this->_loc++;
return Token{TOKEN::COMMA};
case '"': {
size_t start = this->_loc;
this->_loc++;
while (_data[this->_loc] != '"' && this->hasMoreTokens()) {
this->_loc++;
}
size_t start = ++this->_loc;
while (_data[++this->_loc] != '"' && this->hasMoreTokens())
;
if (!this->hasMoreTokens()) {
throw new std::runtime_error("no more chars to parse");
}
// pass the last "
this->_loc++;
return Token{TOKEN::STRING, start, this->_loc};
return Token{TOKEN::STRING, start, ++this->_loc};
}
default: {
size_t start = this->_loc;
@ -63,66 +60,105 @@ Token Tokenizer::getToken() {
}
}
// JSONNode parseJson(char* data, size_t len) {
// auto tokenizer = Tokenizer(data, len);
// std::stack<JSONNode*> ctx;
// auto token = tokenizer.getToken();
// auto parent = new JSONNode();
JSONNode parseJson(const char* data, size_t len) {
auto tokenizer = Tokenizer(data, len);
std::stack<JSONNode*> ctx;
while (tokenizer.hasMoreTokens()) {
auto token = tokenizer.getToken();
switch (token.type) {
case TOKEN::NULL_TYPE:
break;
case TOKEN::CURLY_OPEN:
ctx.push(new JSONNode(JSONNode::Type::OBJECT));
break;
case TOKEN::CURLY_CLOSE:
if (ctx.top()->is_type(JSONNode::Type::OBJECT)) {
auto obj = ctx.top();
ctx.pop();
// top-level object
if (ctx.size() == 0) {
return *obj;
}
// switch (token.type) {
// case TOKEN::CURLY_OPEN:
// parent->setObject(new JSONObject());
// break;
// case TOKEN::ARRAY_OPEN:
// parent->setList(new JSONList());
// break;
// default:
// throw new std::runtime_error("invalid json");
// }
// ctx.push(parent);
if (ctx.top()->is_type(JSONNode::Type::STRING)) {
auto key = ctx.top();
ctx.pop();
if (ctx.top()->is_type(JSONNode::Type::OBJECT)) {
ctx.top()->getObject()->insert({key->getString(), obj});
} else {
throw new std::runtime_error("invalid json");
}
} else if (ctx.top()->is_type(JSONNode::Type::LIST)) {
auto list = ctx.top()->getList();
list->push_back(obj);
}
}
break;
case TOKEN::ARRAY_OPEN:
ctx.push(new JSONNode(JSONNode::Type::LIST));
break;
case TOKEN::ARRAY_CLOSE:
if (ctx.top()->is_type(JSONNode::Type::STRING)) {
// key is above
auto key = ctx.top();
ctx.pop();
if (ctx.top()->is_type(JSONNode::Type::OBJECT)) {
ctx.top()->getObject()->insert({key->getString(), new JSONNode()});
}
} else if (ctx.top()->is_type(JSONNode::Type::LIST)) {
auto obj = ctx.top();
ctx.pop();
// top-level object
if (ctx.size() == 0) {
return *obj;
}
// while (tokenizer.hasMoreTokens()) {
// auto token = tokenizer.getToken();
// switch (token.type) {
// case TOKEN::CURLY_OPEN:
// ctx.push(new JSONNode(JSONNode::Type::OBJECT));
// break;
// case TOKEN::CURLY_CLOSE:
// if (ctx.top()->is_type(JSONNode::Type::OBJECT)) {
// auto obj = ctx.top();
// ctx.pop();
// if (ctx.top()->is_type(JSONNode::Type::LIST)) {
// auto list = ctx.top()->getList();
// list->push_back(obj);
// } else if (ctx.top()->is_type(JSONNode::Type::STRING)) {
// //
// auto key = ctx.top();
// ctx.pop();
// if (ctx.top()->is_type(JSONNode::Type::OBJECT)) {
// ctx.top()->getObject()->insert({key->getString(), obj});
// }
// }
// } else {
// throw new std::runtime_error("invalid json");
// }
// break;
// case TOKEN::COLON:
// break;
// case TOKEN::ARRAY_OPEN:
// break;
// case TOKEN::ARRAY_CLOSE:
// break;
// case TOKEN::COMMA:
// break;
// case TOKEN::NULL_TYPE:
// break;
// case TOKEN::STRING:
// break;
// case TOKEN::NUMBER:
// break;
// }
// }
// }
if (ctx.top()->is_type(JSONNode::Type::LIST)) {
auto list = ctx.top()->getList();
list->push_back(obj);
}
}
break;
case TOKEN::STRING: {
auto str =
new std::string(data + token.start, token.end - token.start - 1);
if (ctx.top()->is_type(JSONNode::Type::OBJECT)) {
ctx.push(new JSONNode(str));
} else if (ctx.top()->is_type(JSONNode::Type::STRING)) {
auto key = ctx.top();
ctx.pop();
if (ctx.top()->is_type(JSONNode::Type::OBJECT)) {
ctx.top()->getObject()->insert(
{key->getString(), new JSONNode(str)});
} else {
throw new std::runtime_error("invalid json");
}
} else if (ctx.top()->is_type(JSONNode::Type::LIST)) {
ctx.top()->getList()->push_back(new JSONNode(str));
}
break;
}
case TOKEN::NUMBER:
if (ctx.top()->is_type(JSONNode::Type::LIST)) {
ctx.top()->getList()->push_back(new JSONNode(JSONNode::Type::NUMBER));
} else if (ctx.top()->is_type(JSONNode::Type::STRING)) {
auto key = ctx.top();
ctx.pop();
if (ctx.top()->is_type(JSONNode::Type::OBJECT)) {
ctx.top()->getObject()->insert(
{key->getString(), new JSONNode(JSONNode::Type::NUMBER)});
} else {
throw new std::runtime_error("invalid json");
}
}
break;
case TOKEN::COMMA:
break;
case TOKEN::COLON:
break;
}
}
}
} // namespace io
@ -141,7 +177,7 @@ std::map<std::string, array> load_safetensor(
in_stream->read(reinterpret_cast<char*>(&jsonHeaderLength), 8);
if (jsonHeaderLength <= 0) {
throw std::runtime_error(
"[load_safetensor] Invalid json header lenght " + in_stream->label());
"[load_safetensor] Invalid json header length " + in_stream->label());
}
// Load the json metadata
char json[jsonHeaderLength];

34
mlx/safetensor.h
View File

@ -13,8 +13,8 @@ namespace mlx::core {
namespace io {
class JSONNode;
using JSONObject = std::map<std::string, std::shared_ptr<JSONNode>>;
using JSONList = std::vector<std::shared_ptr<JSONNode>>;
using JSONObject = std::map<std::string, JSONNode*>;
using JSONList = std::vector<JSONNode*>;
class JSONNode {
public:
@ -29,11 +29,39 @@ class JSONNode {
this->_values.list = new JSONList();
}
};
JSONNode(std::string* s) : _type(Type::STRING) {
this->_values.s = s;
};
JSONObject* getObject() {
if (!is_type(Type::OBJECT)) {
throw new std::runtime_error("not an object");
}
return this->_values.object;
}
JSONList* getList() {
if (!is_type(Type::LIST)) {
throw new std::runtime_error("not a list");
}
return this->_values.list;
}
std::string getString() {
if (!is_type(Type::STRING)) {
throw new std::runtime_error("not a string");
}
return *this->_values.s;
}
inline bool is_type(Type t) {
return this->_type == t;
}
inline Type type() const {
return this->_type;
}
private:
union Values {
JSONObject* object;
@ -44,6 +72,8 @@ class JSONNode {
Type _type;
};
JSONNode parseJson(const char* data, size_t len);
enum class TOKEN {
CURLY_OPEN,
CURLY_CLOSE,

43
tests/load_tests.cpp
View File

@ -27,11 +27,48 @@ TEST_CASE("test tokenizer") {
CHECK_EQ(tokenizer.getToken().type, io::TOKEN::ARRAY_CLOSE);
CHECK_EQ(tokenizer.getToken().type, io::TOKEN::CURLY_CLOSE);
CHECK_EQ(tokenizer.getToken().type, io::TOKEN::NULL_TYPE);
raw = std::string(" { \"testing\": \"test\"} ");
tokenizer = io::Tokenizer(raw.c_str(), raw.size());
CHECK_EQ(tokenizer.getToken().type, io::TOKEN::CURLY_OPEN);
CHECK_EQ(tokenizer.getToken().type, io::TOKEN::STRING);
CHECK_EQ(tokenizer.getToken().type, io::TOKEN::COLON);
CHECK_EQ(tokenizer.getToken().type, io::TOKEN::STRING);
CHECK_EQ(tokenizer.getToken().type, io::TOKEN::CURLY_CLOSE);
CHECK_EQ(tokenizer.getToken().type, io::TOKEN::NULL_TYPE);
}
// TEST_CASE("test load_safetensor") {
// auto array = load_safetensor("../../temp.safe");
// }
TEST_CASE("test load_safetensor") {
auto raw = std::string("{}");
auto res = io::parseJson(raw.c_str(), raw.size());
CHECK(res.is_type(io::JSONNode::Type::OBJECT));
raw = std::string("[]");
res = io::parseJson(raw.c_str(), raw.size());
CHECK(res.is_type(io::JSONNode::Type::LIST));
raw = std::string("[{}, \"test\"]");
res = io::parseJson(raw.c_str(), raw.size());
CHECK(res.is_type(io::JSONNode::Type::LIST));
CHECK_EQ(res.getList()->size(), 2);
CHECK(res.getList()->at(0)->is_type(io::JSONNode::Type::OBJECT));
CHECK(res.getList()->at(1)->is_type(io::JSONNode::Type::STRING));
MESSAGE(res.getList()->at(1)->getString());
CHECK_EQ(res.getList()->at(1)->getString(), "test");
raw = std::string("{\"test\": \"test\", \"test_num\": 1}");
res = io::parseJson(raw.c_str(), raw.size());
CHECK(res.is_type(io::JSONNode::Type::OBJECT));
CHECK_EQ(res.getObject()->size(), 2);
CHECK(res.getObject()->at("test")->is_type(io::JSONNode::Type::STRING));
CHECK_EQ(res.getObject()->at("test")->getString(), "test");
CHECK(res.getObject()->at("test_num")->is_type(io::JSONNode::Type::NUMBER));
raw = std::string("{\"test\": { \"test\": \"test\"}}");
res = io::parseJson(raw.c_str(), raw.size());
CHECK(res.is_type(io::JSONNode::Type::OBJECT));
CHECK_EQ(res.getObject()->size(), 1);
CHECK(res.getObject()->at("test")->is_type(io::JSONNode::Type::OBJECT));
CHECK_EQ(res.getObject()->at("test")->getObject()->size(), 1);
CHECK(res.getObject()->at("test")->getObject()->at("test")->is_type(
io::JSONNode::Type::STRING));
}
TEST_CASE("test single array serialization") {
// Basic test