cli.

Author: trivialfis

src/cli_main.cc

@@ -329,29 +329,45 @@ class CLI {
   }
 
   void LoadModel(std::string const& path, Learner* learner) const {
-    if (common::FileExtension(path) == "json") {
-      auto buffer = common::LoadSequentialFile(path);
-      CHECK_GT(buffer.size(), 2);
-      CHECK_EQ(buffer[0], '{');
-      Json in{Json::Load({buffer.data(), buffer.size()})};
+    auto ext = common::FileExtension(path);
+    auto read_file = [&]() {
+      auto str = common::LoadSequentialFile(path);
+      CHECK_GE(str.size(), 3);  // "{}\0"
+      CHECK_EQ(str[0], '{');
+      return str;
+    };
+
+    if (ext == "json") {
+      auto buffer = read_file();
+      Json in{Json::Load(StringView{buffer.data(), buffer.size()})};
+      learner->LoadModel(in);
+    } else if (ext == "ubj") {
+      auto buffer = read_file();
+      Json in = Json::Load(StringView{buffer.data(), buffer.size()}, std::ios::binary);
       learner->LoadModel(in);
     } else {
-      std::unique_ptr<dmlc::Stream> fi(dmlc::Stream::Create(path.c_str(), "r"));
-      learner->LoadModel(fi.get());
+      LOG(FATAL) << "Unknown model format:" << path << ", expecting either json or ubj.";
     }
   }
 
   void SaveModel(std::string const& path, Learner* learner) const {
     learner->Configure();
     std::unique_ptr<dmlc::Stream> fo(dmlc::Stream::Create(path.c_str(), "w"));
-    if (common::FileExtension(path) == "json") {
+    auto ext = common::FileExtension(path);
+    auto save_json = [&](std::ios::openmode mode) {
       Json out{Object()};
       learner->SaveModel(&out);
-      std::string str;
-      Json::Dump(out, &str);
-      fo->Write(str.c_str(), str.size());
+      std::vector<char> str;
+      Json::Dump(out, &str, mode);
+      fo->Write(str.data(), str.size());
+    };
+
+    if (ext == "json") {
+      save_json(std::ios::out);
+    } else if (ext == "ubj") {
+      save_json(std::ios::binary);
     } else {
-      learner->SaveModel(fo.get());
+      LOG(FATAL) << "Unknown model format:" << path << ", expecting either json or ubj.";
     }
   }