wehub-resource-sync
285 lines
10 KiB
C++
285 lines
10 KiB
C++
/*!
|
|
* \file support/json_parser.h
|
|
* \brief Helps to parse JSON strings and objects.
|
|
*/
|
|
#ifndef MLC_LLM_SUPPORT_JSON_PARSER_H_
|
|
#define MLC_LLM_SUPPORT_JSON_PARSER_H_
|
|
|
|
#include <tvm/ffi/container/shape.h>
|
|
#include <tvm/ffi/dtype.h>
|
|
#include <tvm/ffi/extra/json.h>
|
|
#include <tvm/ffi/string.h>
|
|
#include <tvm/runtime/logging.h>
|
|
|
|
#include <optional>
|
|
|
|
#include "result.h"
|
|
|
|
namespace mlc {
|
|
namespace llm {
|
|
namespace json {
|
|
|
|
using ::tvm::ffi::json::Array;
|
|
using ::tvm::ffi::json::Object;
|
|
using ::tvm::ffi::json::Value;
|
|
|
|
/*!
|
|
* \brief Parse a JSON string to a JSON object.
|
|
* \param json_str The JSON string to parse.
|
|
* \return The parsed JSON object.
|
|
*/
|
|
inline Object ParseToJSONObject(const std::string& json_str) {
|
|
tvm::ffi::String err;
|
|
Value result = ::tvm::ffi::json::Parse(json_str, &err);
|
|
TVM_FFI_CHECK(err.empty(), ValueError)
|
|
<< "Failed to parse JSON: err. The JSON string is:" << json_str;
|
|
auto opt = result.try_cast<Object>();
|
|
TVM_FFI_CHECK(opt.has_value(), ValueError)
|
|
<< "The given string is not a JSON object: " << json_str;
|
|
return *opt;
|
|
}
|
|
/*!
|
|
* \brief Parse a JSON string to a JSON object.
|
|
* \param json_str The JSON string to parse.
|
|
* \return The parsed JSON object, or the error message.
|
|
*/
|
|
inline Result<Object> ParseToJSONObjectWithResultReturn(const std::string& json_str) {
|
|
using TResult = Result<Object>;
|
|
tvm::ffi::String err;
|
|
Value result = ::tvm::ffi::json::Parse(json_str, &err);
|
|
if (!err.empty()) {
|
|
return TResult::Error("Failed to parse JSON: err. The JSON string is: " + json_str +
|
|
". The error is " + std::string(err));
|
|
}
|
|
auto opt = result.try_cast<Object>();
|
|
if (!opt.has_value()) {
|
|
return TResult::Error("ValueError: The given string is not a JSON object: " + json_str);
|
|
}
|
|
return TResult::Ok(*opt);
|
|
}
|
|
|
|
/*!
|
|
* \brief Lookup a JSON object by a key, and convert it to a given type.
|
|
* \param json The JSON object to look up.
|
|
* \param key The key to look up.
|
|
* \tparam ValueType The type to be converted to.
|
|
* \return The converted value.
|
|
*/
|
|
template <typename ValueType>
|
|
ValueType Lookup(const Object& json, const std::string& key);
|
|
/*!
|
|
* \brief Lookup a JSON array by an index, and convert it to a given type.
|
|
* \param json The JSON array to look up.
|
|
* \param index The index to look up.
|
|
* \tparam ValueType The type to be converted to.
|
|
* \return The converted value.
|
|
*/
|
|
template <typename ValueType>
|
|
ValueType Lookup(const Array& json, int index);
|
|
/*!
|
|
* \brief Lookup a JSON object by a key, and convert it to a given type.
|
|
* If the key doesn't exist or has null value, the default value is returned.
|
|
* \param json The JSON object to look up.
|
|
* \param key The key to look up.
|
|
* \tparam ValueType The type to be converted to.
|
|
* \return The converted value, or the default value if the key doesn't exist or has null value.
|
|
*/
|
|
template <typename ValueType>
|
|
inline ValueType LookupOrDefault(const Object& json, const std::string& key,
|
|
const ValueType& default_value) {
|
|
if (json.count(key) == 0 || json.at(key) == nullptr) {
|
|
return default_value;
|
|
}
|
|
auto opt = json.at(key).try_cast<ValueType>();
|
|
TVM_FFI_CHECK(opt.has_value(), ValueError) << "key `" << key << "` has unexpected type";
|
|
return *opt;
|
|
}
|
|
/*!
|
|
* \brief Lookup a JSON object by a key, and convert it to a given type.
|
|
* If the key doesn't exist or has null value, return std::nullopt.
|
|
* \param json The JSON object to look up.
|
|
* \param key The key to look up.
|
|
* \tparam ValueType The type to be converted to.
|
|
* \return The converted value, or std::nullopt if the value doesn't exist or has null value.
|
|
*/
|
|
template <typename ValueType>
|
|
inline std::optional<ValueType> LookupOptional(const Object& json, const std::string& key) {
|
|
if (json.count(key) == 0 || json.at(key) == nullptr) {
|
|
return std::nullopt;
|
|
}
|
|
auto opt = json.at(key).try_cast<ValueType>();
|
|
TVM_FFI_CHECK(opt.has_value(), ValueError) << "key `" << key << "` has unexpected type";
|
|
return *opt;
|
|
}
|
|
/*!
|
|
* \brief Lookup a JSON object by a key, and convert it to a given type.
|
|
* \param json The JSON object to look up.
|
|
* \param key The key to look up.
|
|
* \tparam ValueType The type to be converted to.
|
|
* \return The converted value, or the error message.
|
|
*/
|
|
template <typename ValueType>
|
|
inline Result<ValueType> LookupWithResultReturn(const Object& json, const std::string& key) {
|
|
using TResult = Result<ValueType>;
|
|
if (json.count(key) == 0) {
|
|
return TResult::Error("ValueError: key \"" + key + "\" not found in the JSON object");
|
|
}
|
|
auto opt = json.at(key).try_cast<ValueType>();
|
|
if (!opt.has_value()) {
|
|
return TResult::Error("ValueError: key \"" + key + "\" has unexpected value type.");
|
|
}
|
|
return TResult::Ok(*opt);
|
|
}
|
|
/*!
|
|
* \brief Lookup a JSON object by a key, and convert it to a given type.
|
|
* If the key doesn't exist or has null value, the default value is returned.
|
|
* \param json The JSON object to look up.
|
|
* \param key The key to look up.
|
|
* \tparam ValueType The type to be converted to.
|
|
* \return The converted value, or the default value if the key doesn't exist or has null value
|
|
* , or the error message.
|
|
*/
|
|
template <typename ValueType>
|
|
inline Result<ValueType> LookupOrDefaultWithResultReturn(const Object& json, const std::string& key,
|
|
const ValueType& default_value) {
|
|
using TResult = Result<ValueType>;
|
|
if (json.count(key) == 0 || json.at(key) == nullptr) {
|
|
return TResult::Ok(default_value);
|
|
}
|
|
auto opt = json.at(key).try_cast<ValueType>();
|
|
if (!opt.has_value()) {
|
|
return TResult::Error("ValueError: key \"" + key + "\" has unexpected value type.");
|
|
}
|
|
return TResult::Ok(*opt);
|
|
}
|
|
/*!
|
|
* \brief Lookup a JSON object by a key, and convert it to a given type.
|
|
* If the key doesn't exist or has null value, return std::nullopt.
|
|
* \param json The JSON object to look up.
|
|
* \param key The key to look up.
|
|
* \tparam ValueType The type to be converted to.
|
|
* \return The converted value, or std::nullopt if the value doesn't exist or has null value,
|
|
* , or the error message.
|
|
*/
|
|
template <typename ValueType>
|
|
inline Result<std::optional<ValueType>> LookupOptionalWithResultReturn(const Object& json,
|
|
const std::string& key) {
|
|
using TResult = Result<std::optional<ValueType>>;
|
|
if (json.count(key) == 0 || json.at(key) == nullptr) {
|
|
return TResult::Ok(std::nullopt);
|
|
}
|
|
auto opt = json.at(key).try_cast<ValueType>();
|
|
if (!opt.has_value()) {
|
|
return TResult::Error("ValueError: key \"" + key + "\" has unexpected value type.");
|
|
}
|
|
return TResult::Ok(*opt);
|
|
}
|
|
|
|
// Implementation details
|
|
|
|
/*! \brief Shape extension to incorporate symbolic shapes. */
|
|
struct SymShapeTuple {
|
|
tvm::ffi::Shape shape_values;
|
|
std::vector<std::string> sym_names;
|
|
|
|
/*! \brief Convert symbolic shape tuple to static shape tuple with model config. */
|
|
tvm::ffi::Shape ToStatic(const Object& model_config) {
|
|
std::vector<int64_t> shape;
|
|
shape.reserve(shape_values.size());
|
|
for (int i = 0; i < static_cast<int>(shape_values.size()); ++i) {
|
|
if (shape_values[i] != -1) {
|
|
shape.push_back(shape_values[i]);
|
|
} else {
|
|
auto opt = model_config.at(sym_names[i]).try_cast<int64_t>();
|
|
TVM_FFI_CHECK(opt.has_value(), ValueError)
|
|
<< "model config is expected to contain \"" << sym_names[i]
|
|
<< "\" as an integer. However, the given config has unexpected type for \""
|
|
<< sym_names[i] << "\".";
|
|
shape.push_back(*opt);
|
|
}
|
|
}
|
|
return tvm::ffi::Shape(std::move(shape));
|
|
}
|
|
};
|
|
|
|
namespace details {
|
|
|
|
inline DLDataType DTypeFromString(const std::string& s) { return tvm::ffi::StringToDLDataType(s); }
|
|
|
|
inline SymShapeTuple SymShapeTupleFromArray(const Array& shape) {
|
|
std::vector<int64_t> result;
|
|
std::vector<std::string> sym_names;
|
|
result.reserve(shape.size());
|
|
sym_names.reserve(shape.size());
|
|
for (int i = 0; i < static_cast<int>(shape.size()); ++i) {
|
|
const auto& dim = shape[i];
|
|
auto str_opt = dim.try_cast<std::string>();
|
|
if (str_opt.has_value()) {
|
|
result.push_back(-1);
|
|
sym_names.push_back(*str_opt);
|
|
} else {
|
|
auto int_opt = dim.try_cast<int64_t>();
|
|
TVM_FFI_CHECK(int_opt.has_value(), ValueError) << "shape has unexpected type";
|
|
result.push_back(*int_opt);
|
|
sym_names.push_back("");
|
|
}
|
|
}
|
|
return SymShapeTuple{tvm::ffi::Shape(std::move(result)), sym_names};
|
|
}
|
|
|
|
} // namespace details
|
|
|
|
template <typename ValueType>
|
|
inline ValueType Lookup(const Object& json, const std::string& key) {
|
|
TVM_FFI_CHECK(json.count(key) != 0, ValueError)
|
|
<< "key `" << key << "` not found in the JSON object";
|
|
auto opt = json.at(key).try_cast<ValueType>();
|
|
TVM_FFI_CHECK(opt.has_value(), ValueError) << "key `" << key << "` has unexpected type";
|
|
return *opt;
|
|
}
|
|
|
|
template <typename ValueType>
|
|
inline ValueType Lookup(const Array& json, int index) {
|
|
TVM_FFI_ICHECK(index < static_cast<int>(json.size()))
|
|
<< "IndexError: json::array index out of range";
|
|
auto opt = json[index].try_cast<ValueType>();
|
|
TVM_FFI_ICHECK(opt.has_value()) << "ValueError: value at index `" << index
|
|
<< "` has unexpected type";
|
|
return *opt;
|
|
}
|
|
|
|
template <>
|
|
inline DLDataType Lookup(const Object& json, const std::string& key) {
|
|
return details::DTypeFromString(Lookup<std::string>(json, key));
|
|
}
|
|
|
|
template <>
|
|
inline DLDataType Lookup(const Array& json, int index) {
|
|
return details::DTypeFromString(Lookup<std::string>(json, index));
|
|
}
|
|
|
|
template <>
|
|
inline SymShapeTuple Lookup(const Object& json, const std::string& key) {
|
|
return details::SymShapeTupleFromArray(Lookup<Array>(json, key));
|
|
}
|
|
|
|
template <>
|
|
inline SymShapeTuple LookupOrDefault(const Object& json, const std::string& key,
|
|
const SymShapeTuple& default_value) {
|
|
if (json.count(key) == 0 || json.at(key) == nullptr) {
|
|
return default_value;
|
|
}
|
|
return details::SymShapeTupleFromArray(Lookup<Array>(json, key));
|
|
}
|
|
|
|
template <>
|
|
inline SymShapeTuple Lookup(const Array& json, int index) {
|
|
return details::SymShapeTupleFromArray(Lookup<Array>(json, index));
|
|
}
|
|
|
|
} // namespace json
|
|
} // namespace llm
|
|
} // namespace mlc
|
|
|
|
#endif // MLC_LLM_SUPPORT_JSON_PARSER_H_
|